From 9555bf311d4e7b2fb02ad24e1374bf87535d1707 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Mon, 19 Oct 2026 17:54:16 +0200 Subject: [PATCH 1/4] working on feeding doubleMM to elbo_site --- src/DoubleMM/f_doubleMM.jl | 23 +- src/HVIApproximation.jl | 23 +- src/PBMApplicator.jl | 55 ++- src/elbo_site.jl | 6 +- test/test_elbo_site_doubleMM.jl | 650 ++++++++++++++++++++++++++++++++ 5 files changed, 717 insertions(+), 40 deletions(-) create mode 100644 test/test_elbo_site_doubleMM.jl diff --git a/src/DoubleMM/f_doubleMM.jl b/src/DoubleMM/f_doubleMM.jl index 82cadc2..4cd3d7e 100644 --- a/src/DoubleMM/f_doubleMM.jl +++ b/src/DoubleMM/f_doubleMM.jl @@ -152,6 +152,17 @@ function f_doubleMM_sites(θc_tr::CA.ComponentMatrix, xPc::CA.ComponentMatrix) #(rep_fac .* r0') .+ (rep_fac .* r1') .* S1 ./ ((rep_fac .* K1') .+ S1) .* S2 ./ ((rep_fac .* K2') .+ S2) end +function f_doubleMM_pop!(pred, θs, xPc::CA.ComponentVector{T}) where T + S1 = view(xPc, Val(:S1), :) + S2 = view(xPc, Val(:S2), :) + is_valid = isfinite.(S1) .&& isfinite.(S2) + @row_valviews r0, r1, K1, K2 = θs + # each variable is a matrix (n_obs x n_site) + pred .= r0' .+ r1' .* S1 ./ (K1' .+ S1) .* S2 ./ (K2' .+ S2) + pred +end + + # function f_doubleMM_sites(θc_tr::CA.ComponentMatrix, xPc::CA.ComponentMatrix) # # extract the parameters as vectors # VT = typeof(CA.getdata(θc_tr)[:,1]) # workaround for non-type-stable Symbol-indexing @@ -292,6 +303,9 @@ function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::Val{scen}) xPvec = int_xP1(vcat(xP_S1, xP_S2)) if (:useSitePBM ∈ scen) PBMSiteApplicator(f_doubleMM; pt.θP, pt.θM, θFix, xPvec) + elseif (:usePopPBM ∈ scen) + local n_MC = 4 + PBMPopulationApplicator(f_doubleMM_pop!, n_MC; pt.θP, pt.θM, θFix, xPvec) else n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) PBMPopulationApplicatorTr(f_doubleMM_sites, n_batch; pt.θP, pt.θM, θFix, xPvec) @@ -404,13 +418,18 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase; θMs_true = int_θMs_sites(transM_sites(ζMs_true)) f_batch = get_hybridproblem_PBmodel(prob; scenario) f = create_nsite_applicator(f_batch, n_site) + f1 = create_nsite_applicator(f_batch, 1) #xP = fill((; S1 = xP_S1, S2 = xP_S2), n_siteall) int_xP_sites = ComponentArrayInterpreter(int_xP1, (n_site,)) xP = int_xP_sites(vcat(repeat(xP_S1, 1, n_site), repeat(xP_S2, 1, n_site))) #xP[:S1,:] #θP = get_θP(prob) # for DoubleMMCase par_templates gives correct θP θP = get_hybridproblem_θP(prob; scenario) - (y_true, _) = f(θP, θMs_true', xP) + n_obs = 8 #TODO? + y_true = zeros(8, n_site) + for i_site in 1:n_site + HVI.apply_model!(y_true[:,i_site], f1, reshape(θP,length(θP),1), θMs_true[:,i_site:i_site], xP[:,i_site]) + end σ_o = FloatType(0.01) #σ_o = FloatType(0.002) logσ2_o = FloatType(2) .* log.(σ_o) @@ -471,6 +490,8 @@ function HVI.get_hybridproblem_HVIApproximation(prob::DoubleMMCase; scenario::Va MeanScalingHVIApproximation(block_ends, FT(2) .* log.(FT(0.1) .* θM[block_ends])) elseif (:sepvar ∈ scen) MeanVarSepHVIApproximation() + elseif (:uniscaling ∈ scen) + MeanUniScalingHVIApproximation(log(0.1)) else MeanHVIApproximationMat() end diff --git a/src/HVIApproximation.jl b/src/HVIApproximation.jl index 76cf1ef..eace3d2 100644 --- a/src/HVIApproximation.jl +++ b/src/HVIApproximation.jl @@ -9,12 +9,12 @@ returns the number of required outputs of the machine learning model per site for a parameter vector `θM`. Subtypes must implement method -`get_numberof_θM(::AbstractHVIApproximation, ml_pred::AbstractArray)` that -returns the number compponents mean parameters `θM` given the machine learning -model predictions, if it differs from the default implementation of -returning their number of rows. -For example, [`MeanScalingHVIApproximation`](@ref) there are more outputs than -the mean parameters with the ML model. +- `get_numberof_θM(::AbstractHVIApproximation, ml_pred::AbstractArray)`: + returns the number of mean parameters in `θM` given the machine learning + model predictions, if it differs from the default implementation of + returning their number of rows. + For example, [`MeanScalingHVIApproximation`](@ref) there are more outputs than + the mean parameters with the ML model. """ abstract type AbstractHVIApproximation end, function get_numberof_MLinputs end @@ -52,6 +52,17 @@ ML model predicts a single scaling factor uniform for all variance parameters struct MeanUniScalingHVIApproximation{TF} <: AbstractMeanScalingHVIApproximation logσ_ζM_base::TF end +function get_numberof_MLinputs(::MeanUniScalingHVIApproximation, θM) + length(θM) + 1 +end +function get_numberof_θM(::MeanUniScalingHVIApproximation, ml_pred::AbstractArray) + # scalingblocks_ends reports the end position of the parameters, hence, the last + # corresponds to the number of overall parameters. + length(ml_pred) - 1 +end + + + """ MeanScalingHVIApproximation(scalingblocks_ends, logσ2_ζM_base) diff --git a/src/PBMApplicator.jl b/src/PBMApplicator.jl index df3394a..2b4bd56 100644 --- a/src/PBMApplicator.jl +++ b/src/PBMApplicator.jl @@ -1,9 +1,9 @@ -struct PBMPopulationApplicator{AXS, MFT, MFTD, IXT, F} <: AbstractPBMApplicator +struct PBMPopulationApplicator{AXS, AXP, MFT, MFTD, F} <: AbstractPBMApplicator fθpop!::F θFixm::MFT θFixm_dc::MFTD axθcomp::Val{AXS} # 3-tuple of axes of θP, θM and θFix - int_xP::IXT + ax_xP::Val{AXP} end # let fmap not descend into isP, because indexing with isP on cpu is faster @@ -30,43 +30,34 @@ can be column-indexed by Symbols. - `xPvec`: `ComponentVector` template of model drivers for a single site """ function PBMPopulationApplicator(fθpop!, n_MC:: Integer; - θP::CA.ComponentVector, θM::CA.ComponentVector, θFix::CA.ComponentVector, + θP::CA.ComponentVector{FT}, θM::CA.ComponentVector, θFix::CA.ComponentVector, xPvec::CA.ComponentVector, diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), - ) where chunk - is_MC = ones(n_MC)' + ) where {FT,chunk} + is_MC = ones(FT, n_MC)' # θPm = CA.ComponentMatrix(CA.getdata(θP) .* is_MC, (CA.getaxes(θP)[1], CA.FlatAxis())) # θMm = CA.ComponentMatrix(CA.getdata(θM) .* is_MC, (CA.getaxes(θM)[1], CA.FlatAxis())) - θFixm = CA.getdata(θFix) .* is_MC + θFixm = convert.(FT,CA.getdata(θFix)) .* is_MC + θFixm_dc = PAT.DiffCache(θFixm, chunk) # θ = vcat(vcat(θPm, θMm), θFixm) axθ = (CA.getaxes(θP)[1], CA.getaxes(θM)[1], CA.getaxes(θFix)[1]) - int_xP = get_concrete(ComponentArrayInterpreter(xPvec)) - # TODO create PAT.DiffCache - θFixm_dc = PAT.DiffCache(θFixm, chunk) - PBMPopulationApplicator(fθpop!, θFixm, θFixm_dc, Val(axθ), int_xP) + ax_xP = CA.getaxes(xPvec)[1] + PBMPopulationApplicator(fθpop!, θFixm, θFixm_dc, Val(axθ), Val(ax_xP)) end get_axis_θP(::PBMPopulationApplicator{AX}) where AX = AX[1] get_axis_θM(::PBMPopulationApplicator{AX}) where AX = AX[2] -function create_nsite_applicator(app::PBMPopulationApplicator, n_site) - error("implenment create_nsite_applicator(app::PBMPopulationApplicator") - # θFix = app.θFixm[1,:] - # isFix = repeat(axes(θFix, 1)', n_site) - # θFixm = if length(θFix) == 0 - # CA.ComponentMatrix(θFix[isFix], (CA.FlatAxis(), CA.FlatAxis())) - # else - # CA.ComponentMatrix(θFix[isFix], (CA.FlatAxis(), CA.getaxes(θFix)[1])) - # end - # # - # intθ = get_concrete(ComponentArrayInterpreter((n_site,), (CA.getaxes(app.intθ)[2],),())) - # int_xP = get_concrete(ComponentArrayInterpreter( - # (), (CA.getaxes(app.int_xP)[1],), (n_site,))) - # rep_fac = ones_similar_x(θFix, n_site) # to reshape into matrix, avoiding repeat - # PBMPopulationApplicator(app.fθpop!, θFixm, rep_fac, intθ, int_xP) +function create_nsite_applicator(app::PBMPopulationApplicator, n_MC; + diffchunk::ForwardDiff.Chunk{chunk} = ForwardDiff.Chunk(8), + ) where chunk + is_MC = ones(eltype(app.θFixm), n_MC)' + θFixm = CA.getdata(app.θFixm[:,1]) .* is_MC + θFixm_dc = PAT.DiffCache(θFixm, chunk) + PBMPopulationApplicator(app.fθpop!, θFixm, θFixm_dc, app.axθcomp, app.ax_xP) end -function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXC}, θsP::AbstractMatrix, θsM::AbstractMatrix, xP) where {AXS,AXC} +function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::AbstractMatrix, θsM::AbstractMatrix, xP) where {AXS,AXP} # function tmpf(θsP, θsM, app) # @allocated Vcat(CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)) # end @@ -86,13 +77,17 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXC}, θsP::Abstrac # AXC, CA.FlatAxis()) # TODO provide PreallocationTools for app.θFixm #θc = VcatCAs(AXS, CA.getdata(θsP), CA.getdata(θsM), CA.getdata(app.θFixm)) - θFixm_dual = PAT.get_tmp(app.θFixm_dc, CA.getdata(θsP)) - if typeof(θFixm_dual) != typeof(app.θFixm) - copyto!(θFixm_dual, app.θFixm) + θFixm_dual = if eltype(θsP) isa ForwardDiff.Dual + copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θsP)), app.θFixm) + else + app.θFixm end θc = VcatCAs(AXS, CA.getdata(θsP), CA.getdata(θsM), θFixm_dual) - xPc = app.int_xP(CA.getdata(xP)) + #xPc = app.int_xP(CA.getdata(xP)) + #xPc = CA.ComponentArray(CA.getdata(xP), AXP, θc.flataxes...) + xPc = CA.ComponentVector(CA.getdata(xP), AXP) app.fθpop!(pred, θc, xPc) + #app.fθpop!(pred, θc, xP) pred end diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 7dde816..53fdc35 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -312,7 +312,8 @@ function exp_nLi!( #θ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) + nL = fLogLik(y_pred, y_o, y_unc) / n_MC + # nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC (; nL,) @@ -364,7 +365,6 @@ function fθpop_dummy2!(pred, θs, xPc::CA.ComponentVector{T}) where T end function fLogLik_dummy(y_pred, y_o, y_unc) - n_obs, n_MC = size(y_pred) - sum(y_pred) / n_MC + sum(y_pred) end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl new file mode 100644 index 0000000..0df7101 --- /dev/null +++ b/test/test_elbo_site_doubleMM.jl @@ -0,0 +1,650 @@ +ENV["MLDATADEVICES_SILENCE_WARN_NO_GPU"]="1" # suppress warning on missing CUDA +#using LinearAlgebra, BlockDiagonals +using LinearAlgebra +using StatsFuns: logistic + +using Test +using HybridVariationalInference +using HybridVariationalInference: HybridVariationalInference as CP +using HybridVariationalInference: HybridVariationalInference as HVI +using StableRNGs +using Random +using ComponentArrays: ComponentArrays as CA +#using TransformVariables +using Bijectors +#import PreallocationTools as PAT +import JLD2 +import PDMats + +import Lux +# import Enzyme +import ForwardDiff + + +rng = StableRNG(1234) + +prob = DoubleMM.DoubleMMCase() +scenario = Val((:use_Lux,:usePopPBM,:covarK2,:uniscaling,)) +diffchunk = ForwardDiff.Chunk(8) + +(; xM, θP_true, θMs_true, xP, y_true, y_o, y_unc,) = gen_hybridproblem_synthetic(rng, prob; scenario); +i_indiv_batch = 10 .+ (1:20) +n_site = length(i_indiv_batch) +par_templates = get_hybridproblem_par_templates(prob; scenario) +n_θP = length(par_templates.θP) +n_θM = length(par_templates.θM) +n_cov = get_hybridproblem_n_covar(prob; scenario) +pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) +n_covP = length(pbm_covars) + + + + g, ϕg = get_hybridproblem_MLapplicator(prob; scenario) + ϕgv = collect(ϕg) + # + n_MC = 3; n_MC_val = Val(3) + n_obs, n_site_all = size(y_o) + transP, transM = get_hybridproblem_transforms(prob; scenario) + # + cor_ends = get_hybridproblem_cor_ends(prob; scenario) + # TODO update init_hybrid_ϕunc + ϕqPc1 = CA.ComponentVector( + μζP = log.(CA.getdata(par_templates.θP)), + logσ_ζP = ones(n_θP) .* log(0.01), + ρsP = [0.1], + ) + ϕqIc1 = CA.ComponentVector( + logσ_ζM = ones(n_θM) .* log(0.01), + ρsM = [0.1], + ) + ϕqP = ϕqP2 = CA.getdata(ϕqPc1) + ϕqI = CA.getdata(ϕqIc1) + intϕqP = get_concrete(ComponentArrayInterpreter(ϕqPc1)) + intϕqI = get_concrete(ComponentArrayInterpreter(ϕqIc1)) + # + indiv_args1 = (;xM = CA.getdata(xM)[:,i_indiv_batch], xP = CA.getdata(xP)[:,i_indiv_batch], + y_o = CA.getdata(y_o)[:,i_indiv_batch], y_unc = CA.getdata(y_unc)[:,i_indiv_batch], + indiv_train = i_indiv_batch) + # + pbm = create_nsite_applicator(get_hybridproblem_PBmodel(prob; scenario), n_MC; diffchunk) + nljoint_args_fix =(; + f! = pbm, + fLogLik = get_hybridproblem_neg_logden_obs(prob; scenario) + ) + + ϕqIcS = CA.ComponentVector(;logσ_ζM_offsets = ϕqIc1.logσ_ζM[2:end] .- ϕqIc1.logσ_ζM[1], ϕqIc1.ρsM) + intϕqIS = get_concrete(ComponentArrayInterpreter(ϕqIcS)) + approxS_ = CP.MeanUniScalingHVIApproximation(ϕqIc1.logσ_ζM[1]) + approxS = get_hybridproblem_HVIApproximation(prob; scenario) + n_M = get_numberof_MLinputs(approxS, par_templates.θM) + h0S = CP.prepare_elbo_helpers(approxS, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); + randn!(h0S.θsP) # for testing should initialized to finite values + +#------------------------------------------------------------- + +@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 "pullback_g_apply!" begin + # @test ϕms0 == ϕms0z + # @test ϕms2 == ϕms2z + # @test size(ϕms0) == (n_M, n_site) + # @test size(ϕms2) == (n_M, n_MC, n_site) + # () -> begin # gradient(sum) + # gr_zygote = Zygote.gradient((ϕg2) -> sum(CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP)), ϕg2v ) + # s, pullback_s_zygote = Zygote.pullback((ϕg2) -> sum(CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP)), ϕg2v ) + # #gr_zygote2 = pullback_s_zygote(ones(eltype(ϕg2v), size(ϕg2v)...)) + # gr_zygote2 = pullback_s_zygote(one(eltype(ϕg2v))) + # @test gr_zygote2[1] ≈ gr_zygote[1] + # y, pullback_zygote = Zygote.pullback((ϕg2) -> CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP), ϕg2v ) + # gr_zygote3 = pullback_zygote(ones(eltype(y), size(y)...)) + # gr_zygote3[1] ≈ gr_zygote[1] + # end + # # + # # concatenate function f3(g(ϕ_g)) + # f3 = (x) -> sum(3.0 .* x) + # CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP) + # # one pass of composed function + # gr_zygote = Zygote.gradient((ϕg2, ζsP) -> f3(CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP)), ϕg2v, ζsP ) + # s, pullback_s_zygote = Zygote.pullback((ϕg2) -> f3(CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP)), ϕg2v ) + # gr_zygote2 = pullback_s_zygote(one(eltype(s))) + # @test gr_zygote2[1] ≈ gr_zygote[1] + # # mixed AD, differentiate f3 by FowardDiff and pull back through g + # y_oop = CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP) + # #gr_h = Zygote.gradient(y -> sum(f3(y)), y1)[1] + # gr_h = ForwardDiff.gradient(y -> sum(f3(y)), y_oop) + # y, pullback_zygote = Zygote.pullback((ϕg2) -> CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP), ϕg2v ) + # @test y == y_oop + # gr_zygote3 = pullback_zygote(gr_h) + # gr_zygote3[1] ≈ gr_zygote[1] + # # + # dϕg = zero(ϕg2v) + # dζsP = zero(ζsP) + # #Enzyme.make_one!(y) + # y .= rand() + # dϕg .= rand() # check that initial values do not effect result + # dζsP .= rand() # check that initial values do not effect result + # dy = convert.(eltype(y), gr_h) + # CP.pullback_g_apply!(y, dϕg, dζsP, dy, ϕg2v, xM, ζsP, pbm_covar_indices2, g2, h) + # @test y == y_oop + # @test dϕg ≈ gr_zygote[1] + # @test dζsP ≈ gr_zygote[2] rtol=1e-3 + # @test dy ≈ gr_h # not modified + # #@benchmark CP.pullback_g_apply!(y, dϕg, dy, ϕg2v, xM, ζsP, pbm_covar_indices2, g2, h) + # # + # () -> begin # explicitly splitting the forward and backward pass + # # they get cached anymay and require allocating the Duplicated Wrappers twice + # # hence there is no performance benefit + # # Compile once outside the hot loop + # fwd, rev = Enzyme.autodiff_thunk( + # Enzyme.ReverseSplitNoPrimal, + # Enzyme.Const{typeof(g_apply!)}, + # Enzyme.Const, + # Enzyme.Duplicated{typeof(y)}, + # Enzyme.Duplicated{typeof(ϕg2v)}, + # Enzyme.Const{typeof(xM)}, + # Enzyme.Const{typeof(ζP)}, + # Enzyme.Const{typeof(pbm_covar_indices2)}, + # Enzyme.Const{typeof(g2)}, + # Enzyme.Duplicated{typeof(h.xMP)} + # ) + # # take care, dy is also modified + # function grad2_g_apply!(y, dϕg, dy, ϕg, xM, ζP, pbm_covar_indices, g, h, fwd, rev) + # fill!(dϕg, zero(eltype(dϕg))) + # fill!(h.dxMP, zero(eltype(h.dxMP))) + # copyto!(h.dy, dy) # copy to avoid modifying dy + # tape, _, _ = fwd( + # Enzyme.Const(g_apply!), + # Enzyme.Duplicated(y, h.dy), + # Enzyme.Duplicated(ϕg, dϕg), + # Enzyme.Const(xM), + # Enzyme.Const(ζP), + # Enzyme.Const(pbm_covar_indices), + # Enzyme.Const(g), + # Enzyme.Duplicated(h.xMP, h.dxMP) + # ) + # rev( + # Enzyme.Const(g_apply!), + # Enzyme.Duplicated(y, h.dy), + # Enzyme.Duplicated(ϕg, dϕg), + # Enzyme.Const(xM), + # Enzyme.Const(ζP), + # Enzyme.Const(pbm_covar_indices), + # Enzyme.Const(g), + # Enzyme.Duplicated(h.xMP, h.dxMP), + # tape + # ) + # return nothing + # end + + # dy = convert.(eltype(y), gr_h) + # grad2_g_apply!(y, dϕg, dy, ϕg2v, xM, ζP, pbm_covar_indices2, g2, h, fwd, rev) + # @test y == y_oop + # @test dϕg ≈ gr_zygote[1] + # @test dy ≈ gr_h # not modified + # #@usingany BenchmarkTools + # #@benchmark grad2_g_apply!(y, dϕg, dy, ϕg2v, xM, ζP, pbm_covar_indices2, g2, h, fwd, rev) + # end +end + +@testset "pullback_sample_ζsP!" begin + # ϕqPc = intϕqP(ϕqP) + # rnormP = zero(ζsP) + # randn!(rnormP) # before input gaussian noise + # ζsP .= 0 + # #logσ_ζP = zero(ϕqPc.logσ_ζP) # cretes a view rather than copy + # logσ_ζP = zero(ϕqPc.logσ_ζP) + # CP.sample_ζsP!(ζsP, logσ_ζP, rnormP, ϕqPc) + # mean(ζsP; dims=2) + # ζsP1 = copy(ζsP) + + # # Enzyme result via the mutating routine (2-D: n_θP * n_MC × n_in) + # dζsP = zero(ζsP) .+ one(eltype(ζsP)) + # dlogσ_ζP = zero(logσ_ζP) .+ one(eltype(ζsP)) + # dϕqc = zero(ϕqPc) + + # randn!(dϕqc) # test that is zerod inside pullback + # dζsP_ = copy(dζsP) + # rnormP_ = copy(rnormP) + # logσ_ζP_ = copy(logσ_ζP) + # dlogσ_ζP_ = copy(dlogσ_ζP) + # #CP.pullback_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, rnormP, logσ_ζP, ϕqPc) # needs rnormP to be noise + # CP.pullback_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, rnormP, ϕqPc) + # @test ζsP == ζsP1 # same forward result + # @test rnormP == rnormP_ + # @test dζsP == dζsP_ + # @test dlogσ_ζP == dlogσ_ζP_ + # @test logσ_ζP == logσ_ζP_ + # # without correlation + # #@test all(dϕqc[Val(:μζP)] .== n_MC) + # # #@test dϕqc[Val(:logσ_ζP)] ≈ vec(sum(rnormP; dims=2)) # + # dϕqc_comb = copy(dϕqc) + + # randn!(ζsP) # test initial not relevant + # pb_sample_ζsP = CP.primal_pullback_sample_ζsP!(ζsP, logσ_ζP, rnormP, ϕqPc) + # @test ζsP == ζsP1 # same forward result + # @test rnormP ≈ rnormP_# computed the forward pass + # @test logσ_ζP == logσ_ζP_ + # dϕqc .= 0.1 # test initial value not relevant + # #dϕqc .= 0.01 # should not influence results + # pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) + # #pb_sample_ζsP(rnormP, logσ_ζP) + # @test rnormP ≈ rnormP_ # did not modify + # @test logσ_ζP == logσ_ζP_ # not modified + # @test dlogσ_ζP == dlogσ_ζP_ # not modified + # @test dζsP == dζsP_ + # #hcat(dϕqc, dϕqc_comb) + # @test CA.getdata(dϕqc) ≈ CA.getdata(dϕqc_comb) + # # + # # test another pullback + # #dζsP .= dζsP * eltype(dζsP)(2) + # pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) + # @test CA.getdata(dϕqc) ≈ CA.getdata(dϕqc_comb) + # # + # # @usingany BenchmarkTools + # # @benchmark pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) +end + +# to be differentiated, return only first entry of result +function _ftmp2(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] +end + +function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by enzyme + # do not use DiffCache here for helpers_sites + h0p = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; + n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP0, n_M, cor_ends, n_obs, + use_diff_cache = Val(false)) + # and store results to compare to hand-crafted mixed AD + dh0p = Enzyme.make_zero(h0p) + @test dh0p !== h0p # real copy rather than reference + # + indiv_args = indiv_args1 + ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) + dϕ = map(zero, ϕ) + sample_args = (; + approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, + transP, transM) + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕ.ϕg) + randn!(rng1, indiv_args.xM) + primal_enz = tmp = _ftmp2(h0p, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + Enzyme.make_zero!(dϕ) + Enzyme.make_zero!(dh0p) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse) , + _ftmp2, + Enzyme.Active, + Enzyme.Duplicated(h0p, dh0p), + Enzyme.Duplicated(ϕ, dϕ), + Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), + Enzyme.Const(sample_args), + Enzyme.Const(indiv_args), + Enzyme.Const(nljoint_args_fix), + ) + () -> begin + dϕg0_enz ≈ dϕ.ϕg + dϕqP0_enz ≈ dϕ.ϕqP + dϕqI0_enz ≈ dϕ.ϕqI + hcat(dϕg0_enz, dϕ.ϕg) + end + dϕg0_enz = copy(dϕ.ϕg) + dϕqP0_enz = copy(dϕ.ϕqP) + dϕqI0_enz = copy(dϕ.ϕqI) + () -> begin + #@usingany JLD2 + #fname = "intermediate/test_enzyme_dphi0.jld2" + fname = "intermediate/test_enzymeT_dphi0.jld2" + mkpath("intermediate") + JLD2.jldsave(fname, false, IOStream; primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz) + primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz = JLD2.load(fname, + "primal_enz", "dϕg0_enz", "dϕqP0_enz", "dϕqI0_enz"); + end + + h2p = CP.prepare_elbo_helpers(approx, ϕg2, ϕqP; + n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs, + use_diff_cache = Val(false)) + dh2p = Enzyme.make_zero(h2p) + indiv_args = indiv_args1 + ϕ2 = (;ϕg=ϕg2v, ϕqP, ϕqI) + dϕ2 = map(zero, ϕ2) + sample_args2 = (; + approx, g=g2, is_testmode, pbm_covar_indices = pbm_covar_indices2, intϕqP, intϕqI, cor_ends, + transP, transM) + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕ2.ϕg) + randn!(rng1, indiv_args.xM) + primal_enz = _ftmp2(h2p, ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix) + Enzyme.make_zero!(dϕ2) + Enzyme.make_zero!(dh2p) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse) , + _ftmp2, + Enzyme.Active, + Enzyme.Duplicated(h2p, dh2p), + Enzyme.Duplicated(ϕ2, dϕ2), + Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), + Enzyme.Const(sample_args2), + Enzyme.Const(indiv_args), + Enzyme.Const(nljoint_args_fix), + ) + () -> begin + primal2_enz ≈ primal_enz + dϕg2_enz ≈ dϕ2.ϕg + @test dϕqP2_enz ≈ dϕ2.ϕqP atol=0.001 + dϕqI2_enz ≈ dϕ2.ϕqI + hcat(dϕg2_enz, dϕ2.ϕg) + hcat(dϕqP2_enz, dϕ2.ϕqP) + end + dϕg2_enz = copy(dϕ2.ϕg) + dϕqP2_enz = copy(dϕ2.ϕqP) + dϕqI2_enz = copy(dϕ2.ϕqI) + () -> begin + #fname = "intermediate/test_enzyme_dphi2.jld2" + fname = "intermediate/test_enzymeT_dphi2.jld2" + mkpath("intermediate") + JLD2.jldsave(fname, false, IOStream; primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz) + primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz = JLD2.load(fname, + "primal2_enz", "dϕg2_enz", "dϕqP2_enz", "dϕqI2_enz"); + end +end + +function grad_neg_elbo_sites_enzyme_scaled() + # now with approxS and respective arguments + is_testmode = false + indiv_args = indiv_args1 + ϕ2S = (;ϕg=ϕg2v, ϕqP, ϕqI=CA.getdata(ϕqIcS)) + dϕ2S = map(zero, ϕ2S) + sample_args2S = (; + approx=approxS, g=g2, is_testmode, pbm_covar_indices = pbm_covar_indices2, + intϕqP, intϕqI=intϕqIS, cor_ends, + transP, transM) + h2Sp = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; + n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, use_diff_cache = Val(false)) + dh2Sp = Enzyme.make_zero(h2Sp) + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕ2S.ϕg) + randn!(rng1, indiv_args.xM) + primal_enz2S = tmp = _ftmp2(h2Sp, ϕ2S, rnormPM, sample_args2S, indiv_args) + Enzyme.make_zero!(dϕ2S) + Enzyme.make_zero!(dh2Sp) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse) , + _ftmp2, + Enzyme.Active, + Enzyme.Duplicated(h2Sp, dh2Sp), + Enzyme.Duplicated(ϕ2S, dϕ2S), + Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), + Enzyme.Const(sample_args2S), + Enzyme.Const(indiv_args), + ) + dϕg2S_enz = copy(dϕ2S.ϕg) + dϕqP2S_enz = copy(dϕ2S.ϕqP) + dϕqI2S_enz = copy(dϕ2S.ϕqI) + () -> begin + #fname = "intermediate/test_enzyme_dphi2.jld2" + fname = "intermediate/test_enzymeS_dphi2.jld2" + mkpath("intermediate") + JLD2.jldsave(fname, false, IOStream; primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz) + primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load(fname, + "primal_enz2S", "dϕg2S_enz", "dϕqP2S_enz", "dϕqI2S_enz"); + end +end + +function tmp_reference_revision() + #@usingany ReferenceRevision + ref = open_process(rev="origin/sitepar", path="tmp/ref", use=[:Revise, :ShareAdd]) + ref.eval(:(Pkg.activate("tmp/ref/test"))) + #ref.eval(:(Pkg.instantiate())) + #ref.Pkg.project().path + ref.eval(:(using HybridVariationalInference)) + ref.eval(:(import HybridVariationalInference as HVI)) + ref.eval(:(using Lux)) # to deserialize objects + ref.eval(:(using SharedArrays)) # to deserialize objects + #ref.eval(:(@usingany Infiltrator)) # not output during infiltrat in revision + tmp = ref.HVI.test3(3) + ref.Revise.retry() # after every change of source file + res0v, gradh0v = ref.HVI.grad_neg_elbo_sites( + h0, + (;), + approx, + rnormPM, + ϕgv, ϕqP, ϕqI, g, nothing; #pbm_covar_indices2; + i_sites_train = 1:n_site, + intϕqP, intϕqI, + xM, + cor_ends, + transP, transM, + is_testmode = false, + ) + close(ref) + + map(k -> k => res0v[k] ≈ res0[k], keys(res0)) + hcat(res0v.dϕqIc, res0.dϕqIc) +end + +@testset "grad_neg_elbo_sites" begin + ϕqIc = intϕqI(ϕqI) + ϕqPc = intϕqP(ϕqP) + is_testmode = false + n_threads_proc = min(Threads.nthreads(), 4) + n_workers = Distributed.nworkers() * n_threads_proc + basesize = n_site ÷ Distributed.nworkers() + distributedEx = Transducers.DistributedEx(;basesize, + threads_basesize = Int(ceil(basesize / n_threads_proc))) + # + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕgv) + randn!(rng1, xM) + indiv_args = indiv_args1 + ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) + sample_args = (; + approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, + transP, transM) + primal = CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + res0, gradh0 = CP.grad_neg_elbo_sites( + #@descend_code_warntype CP.neg_elbo_sites!( + h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + res0_, gradh0_ = CP.grad_neg_elbo_sites( # test deterministic result and distributed + h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, + executor = distributedEx + ) + @test all(map(≈, res0_, res0)) + # # if we saved Enzyme results earlier, compare to them + if isfile("intermediate/test_enzymeT_dphi0.jld2") + primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz = JLD2.load( + "intermediate/test_enzymeT_dphi0.jld2", + "primal_enz", "dϕg0_enz", "dϕqP0_enz", "dϕqI0_enz"); + @test primal_enz ≈ primal[1] + @test dϕg0_enz ≈ res0.dϕg + @test dϕqI0_enz ≈ res0.dϕqIc + @test dϕqP0_enz ≈ res0.dϕqPc + #hcat(dϕqI0_enz, CA.getdata(res0.dϕqIc)) + #hcat(dϕqP0_enz, CA.getdata(res0.dϕqPc)) + #hcat(dϕg0_enz, CA.getdata(res0.dϕg)) + #dϕqP0_enz - CA.getdata(res0.dϕqP) + end + # + #---------------- matrix mode with population covariates + ϕ2 = (;ϕ..., ϕg=ϕg2v) + sample_args2 = (;sample_args..., g=g2, pbm_covar_indices = pbm_covar_indices2) + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕ2.ϕg) + randn!(rng1, indiv_args.xM) + primal2 = CP.neg_elbo_sites!(h2, ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix)[1] + res0, gradh2 = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val) + res0_, gradh2_ = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val, + executor = distributedEx,) + @test all(map(≈, res0_, res0)) + # if we saved Enzyme results earlier, compare to them + if isfile("intermediate/test_enzymeT_dphi2.jld2") + primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz = JLD2.load( + "intermediate/test_enzymeT_dphi2.jld2", + "primal2_enz", "dϕg2_enz", "dϕqP2_enz", "dϕqI2_enz"); + @test primal2_enz ≈ primal2[1] + @test dϕg2_enz ≈ res0.dϕg + @test dϕqI2_enz ≈ res0.dϕqIc + @test dϕqP2_enz ≈ res0.dϕqPc atol = 0.001 + #hcat(dϕqP2_enz, res0.dϕqP) + end + #---------------- approxM with non-empty h.sample_buffers and scaling + # and using eachcol in PBM function + ϕ2S = (;ϕ2..., ϕqI=CA.getdata(ϕqIcS)) + sample_args2S = (;sample_args2..., approx=approxS, intϕqI=intϕqIS) + #nljoint_args_fix2S = nljoint_args_fix =(; f! = pbm_col) + nljoint_args_fix2S = (; nljoint_args_fix..., f! = pbm_col) + # h2Sp = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; + # n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, use_diff_cache = Val(false)) + rng1 = StableRNG(1234) + CP.randnPM!(rng1, rnormPM) + randn!(rng1, ϕ2S.ϕg) + randn!(rng1, indiv_args.xM) + primal2S = CP.neg_elbo_sites!(h2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) + res0, gradh2S = CP.grad_neg_elbo_sites( + h2S, (;), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, + #executor = distributedEx, + ) + if isfile("intermediate/test_enzymeS_dphi2.jld2") + primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load( + "intermediate/test_enzymeS_dphi2.jld2", + "primal_enz2S", "dϕg2S_enz", "dϕqP2S_enz", "dϕqI2S_enz"); + @test primal_enz2S ≈ primal2S[1] + @test dϕg2S_enz ≈ res0.dϕg + @test dϕqI2S_enz ≈ res0.dϕqIc atol=0.01 + @test dϕqP2S_enz ≈ res0.dϕqPc atol=0.01 + #hcat(dϕqP2_enz, res0.dϕqP) + #hcat(dϕqI2S_enz, res0.dϕqI) + end + + function loop_grad_neg_elbo_sites( + n, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.grad_neg_elbo_sites( + h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix; n_MC_val) + end + nothing + end + function alloc_grad_neg_elbo_sites(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + # avoid global variables -> pass them through function + loop_grad_neg_elbo_sites(1, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + #@profview_allocs loop_grad_neg_elbo_sites(1_000,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + @test (@allocated loop_grad_neg_elbo_sites(100,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix)) <= 3_264_000 #3_260_800 + end + alloc_grad_neg_elbo_sites(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) + #_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 + + + + + + From 853ba6d5253f06ef5303abf4d93b49db17a12b82 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Tue, 20 Oct 2026 08:19:43 +0000 Subject: [PATCH 2/4] continue working on feeding doubleMM to elbo_site --- src/DoubleMM/f_doubleMM.jl | 22 +- src/HVIApproximation.jl | 4 + src/ModelApplicator.jl | 16 ++ src/PBMApplicator.jl | 30 ++- src/elbo_site.jl | 5 +- src/init_hybrid_params.jl | 39 +++- src/logden_normal.jl | 17 +- test/test_elbo_site_doubleMM.jl | 384 ++------------------------------ 8 files changed, 133 insertions(+), 384 deletions(-) diff --git a/src/DoubleMM/f_doubleMM.jl b/src/DoubleMM/f_doubleMM.jl index 4cd3d7e..0b073fd 100644 --- a/src/DoubleMM/f_doubleMM.jl +++ b/src/DoubleMM/f_doubleMM.jl @@ -237,7 +237,10 @@ function HVI.get_hybridproblem_MLapplicator( (; transM) = get_hybridproblem_transforms(prob; scenario) lowers, uppers = HVI.get_quantile_transformed(priors, transM) #n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) - range_scaled = if any((:scalingall, :clustered_sites, :clustered_sites2) .∈ Ref(scen)) + approx = get_hybridproblem_HVIApproximation(prob; scenario) + range_scaled = if approx isa AbstractMeanScalingHVIApproximation + #range_scaled = if any((:scalingall, :clustered_sites, :clustered_sites2, :uniscaling) .∈ Ref(scen)) + # TODO make this dependent on approx 1:length(θM) else 1:0 @@ -304,7 +307,7 @@ function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::Val{scen}) if (:useSitePBM ∈ scen) PBMSiteApplicator(f_doubleMM; pt.θP, pt.θM, θFix, xPvec) elseif (:usePopPBM ∈ scen) - local n_MC = 4 + local n_MC = 1 # default n_MC to implement no-MC version of apply_model PBMPopulationApplicator(f_doubleMM_pop!, n_MC; pt.θP, pt.θM, θFix, xPvec) else n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) @@ -338,8 +341,12 @@ end # return pred_global, pred_sites # end -function HVI.get_hybridproblem_neg_logden_obs(::DoubleMMCase; scenario::Val) - neg_logden_indep_normal +function HVI.get_hybridproblem_neg_logden_obs(::DoubleMMCase; scenario::Val{scen}) where scen + if (:usePopPBM ∈ scen) + HVI.neg_logden_indep_normal_sum + else + neg_logden_indep_normal + end end # function HVI.get_hybridproblem_float_type(::DoubleMMCase; scenario) @@ -418,18 +425,13 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase; θMs_true = int_θMs_sites(transM_sites(ζMs_true)) f_batch = get_hybridproblem_PBmodel(prob; scenario) f = create_nsite_applicator(f_batch, n_site) - f1 = create_nsite_applicator(f_batch, 1) #xP = fill((; S1 = xP_S1, S2 = xP_S2), n_siteall) int_xP_sites = ComponentArrayInterpreter(int_xP1, (n_site,)) xP = int_xP_sites(vcat(repeat(xP_S1, 1, n_site), repeat(xP_S2, 1, n_site))) #xP[:S1,:] #θP = get_θP(prob) # for DoubleMMCase par_templates gives correct θP θP = get_hybridproblem_θP(prob; scenario) - n_obs = 8 #TODO? - y_true = zeros(8, n_site) - for i_site in 1:n_site - HVI.apply_model!(y_true[:,i_site], f1, reshape(θP,length(θP),1), θMs_true[:,i_site:i_site], xP[:,i_site]) - end + (y_true, _) = f(θP, θMs_true', xP) σ_o = FloatType(0.01) #σ_o = FloatType(0.002) logσ2_o = FloatType(2) .* log.(σ_o) diff --git a/src/HVIApproximation.jl b/src/HVIApproximation.jl index eace3d2..8a78539 100644 --- a/src/HVIApproximation.jl +++ b/src/HVIApproximation.jl @@ -52,6 +52,10 @@ ML model predicts a single scaling factor uniform for all variance parameters struct MeanUniScalingHVIApproximation{TF} <: AbstractMeanScalingHVIApproximation logσ_ζM_base::TF end +function MeanUniScalingHVIApproximation(app::MeanUniScalingHVIApproximation; + logσ_ζM_base = app.logσ_ζM_base) + MeanUniScalingHVIApproximation(logσ_ζM_base) +end function get_numberof_MLinputs(::MeanUniScalingHVIApproximation, θM) length(θM) + 1 end diff --git a/src/ModelApplicator.jl b/src/ModelApplicator.jl index b6f7d12..34bf5c7 100644 --- a/src/ModelApplicator.jl +++ b/src/ModelApplicator.jl @@ -231,8 +231,24 @@ function apply_model(app::NormalScalingModelApplicator, x, ϕ; kwargs...) # end end +function apply_model!(y, app::NormalScalingModelApplicator, x, ϕ; kwargs...) + apply_model!(y, app.app, x, ϕ; kwargs...) + # @show typeof(app.μ) + # @show typeof(ϕ) + @assert eltype(app.μ) == eltype(ϕ) + if !isempty(app.range_scaled) + view_firstdim(y,app.range_scaled) .= norminvcdf.(app.μ, app.σ, view_firstdim(y,app.range_scaled)) # from StatsFuns + else + y .= norminvcdf.(app.μ, app.σ, y) + end + y +end + + index_firstdim(v::AbstractVector, i) = v[i] index_firstdim(v::AbstractMatrix, i) = v[i,:] +view_firstdim(v::AbstractVector, i) = view(v, i) +view_firstdim(v::AbstractMatrix, i) = view(v,i,:) """ RangeScalingModelApplicator(app, y0) diff --git a/src/PBMApplicator.jl b/src/PBMApplicator.jl index 2b4bd56..dfb8932 100644 --- a/src/PBMApplicator.jl +++ b/src/PBMApplicator.jl @@ -77,7 +77,7 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac # 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) isa ForwardDiff.Dual + θFixm_dual = if eltype(θsP) <: ForwardDiff.Dual copyto!(PAT.get_tmp(app.θFixm_dc, CA.getdata(θsP)), app.θFixm) else app.θFixm @@ -91,5 +91,33 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac pred end +function (app::PBMPopulationApplicator)(θP::AbstractMatrix, θMs_tr::AbstractArray, xP::AbstractMatrix) + 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) + # @assert n_MC = size(θP,2) + pred = zeros(n_obs, n_site, n_MC) + for i_site in 1:n_site + HVI.apply_model!(pred[:,i_site,:], app, θP, θMs_tr[i_site,:,:], xP[:,i_site]) + end + (pred, nothing) +end +function (app::PBMPopulationApplicator)(θP::AbstractVector, θMs_tr::AbstractMatrix, xP::AbstractMatrix) + # for a single MC-run + local n_obs, n_site, pred, θsP, θsM + n_obs = get_n_obs(xP[:,1]) + n_site, n_θM = size(θMs_tr) + θsP = reshape(θP, length(θP), 1) + pred = zeros(n_obs, n_site, 1) + # i_site = 1 + for i_site in 1:n_site + θsM = reshape(view(θMs_tr,i_site,:), n_θM, 1) + apply_model!(view(pred,:,i_site,:), app, θsP, θsM, xP[:,i_site]) + end + (pred[:,:,1], nothing) +end +function get_n_obs(xP::CA.ComponentVector) + isempty(keys(xP)) ? length(xP) : length(xP[keys(xP)[1]]) +end diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 53fdc35..f933cec 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -72,7 +72,6 @@ function compute_nelboi_z!(hi, approx::AbstractHVIApproximation, 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 - (;xP, i_indiv_train) = nljoint_args_indi # 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 @@ -313,8 +312,6 @@ function exp_nLi!( 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( @@ -365,6 +362,6 @@ function fθpop_dummy2!(pred, θs, xPc::CA.ComponentVector{T}) where T end function fLogLik_dummy(y_pred, y_o, y_unc) - sum(y_pred) + sum(y_pred) end diff --git a/src/init_hybrid_params.jl b/src/init_hybrid_params.jl index d5108a1..47ecb9f 100644 --- a/src/init_hybrid_params.jl +++ b/src/init_hybrid_params.jl @@ -183,7 +183,7 @@ end function init_hybrid_ϕunc( approx::SApp, cor_ends::NamedTuple, - ρ0::FT = 0.0f0, + ρ0::FT = zero(FT), logσ2_ζMs::AbstractMatrix{FT} = Array{FT}(undef, 0, 0), logσ2_ζP::AbstractVector{FT} = fill(FT(-10.0), cor_ends.P[end]), ρsP = fill(ρ0, get_cor_count(cor_ends.P)), @@ -191,7 +191,7 @@ function init_hybrid_ϕunc( transM, θM::CA.ComponentVector, n_site::Integer, - relerr = 0.01, + relerr = FT(0.01), ) where {FT, SApp <: MeanScalingHVIApproximation} logσ2 = if isempty(logσ2_ζMs) # relative error of the template of θM @@ -224,3 +224,38 @@ function init_hybrid_ϕunc( (; ϕqc = CA.ComponentVector(;nt...)::CA.ComponentVector, approx) end +# TODO change all σ2 to σ +function init_hybrid_ϕunc( + approx1::MeanUniScalingHVIApproximation, + cor_ends::NamedTuple, + ρ0::FT = 0.0f0, #zero(FT), + logσ2_ζMs::AbstractMatrix{FT} = Array{FT}(undef, 0, 0), + logσ2_ζP::AbstractVector{FT} = fill(FT(-10.0), cor_ends.P[end]), + ρsP = fill(ρ0, get_cor_count(cor_ends.P)), + ρsM = fill(ρ0, get_cor_count(cor_ends.M)); + transM, + θM::CA.ComponentVector, + n_site::Integer, + relerr = FT(0.01), +) where {FT} + local logσ, approx, logσ_ζM_offsets + logσ = if isempty(logσ2_ζMs) + # relative error of the template of θM + σ = compute_σ_unconstrained(transM, CA.getdata(θM), relerr) + #logσ2 = FT(2) * log.(convert.(FT,σ)) + logσ = log.(convert.(FT,σ)) + else + error("check and implement inferring median logσ2 from logσ2_ζMs") + median(logσ2_ζMs; dims=1) + end + # update logσ2_ζM_base of last parameter in approx - its not calibrated + approx = MeanUniScalingHVIApproximation(approx1; logσ_ζM_base = logσ[1]) + logσ_ζM_offsets = logσ[2:end] .- logσ[1] + #tmp = CA.ComponentVector(;zip(Symbol.(axes(is_offset,1)), logσ2_ζM_offsets_gen)...) + nt = (; + logσ_ζP = logσ2_ζP ./ FT(2), + logσ_ζM_offsets, + ρsP, + ρsM) + (; ϕqc = CA.ComponentVector(;nt...)::CA.ComponentVector, approx) +end diff --git a/src/logden_normal.jl b/src/logden_normal.jl index f372e1a..e9b04ff 100644 --- a/src/logden_normal.jl +++ b/src/logden_normal.jl @@ -15,9 +15,9 @@ a low uncertainty estimate and means closer to the observations to help an initial fit. The obtained parameters then can be used as starting values for a the proper fit with `σfac=1.0`. """ -function neg_logden_indep_normal(obs::AbstractArray{TO,D}, μ::AbstractArray{TP,D}, logσ2::AbstractArray{TO}; +function neg_logden_indep_normal(obs::AbstractArray{TO,DO}, μ::AbstractArray{TP,DP}, logσ2::AbstractArray{TO}; i_finobs = .! isnan.(obs), - σfac=one(TO)) where {TO, TP, D} + σfac=one(TO)) where {TO, TP, DO, DP} # log of independent Normal distributions # estimate independent uncertainty of each θM, rather than full covariance #nlogL = sum(σfac .* log.(σs) .+ 1 / 2 .* abs2.((obs .- μ) ./ σs)) @@ -43,9 +43,9 @@ function neg_logden_indep_normal(obs::AbstractArray{TO,D}, μ::AbstractArray{TP, # (σfac .* logσ2_fin .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2_fin)) ./ eltype(μ)(2) nlogL = i_finobs .* (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) - nlogL_sites = if D == 1 + nlogL_sites = if DO == 1 sum(nlogL) # provided a single-site vector - elseif D == 2 + elseif DO == 2 vec(sum(nlogL;dims=1)) else error("need to provide a matrix with observations with one column per site, " * @@ -67,6 +67,15 @@ function neg_logden_indep_normal(obs::AbstractArray{TO,D}, μ::AbstractArray{TP, end +function neg_logden_indep_normal_sum(obs::AbstractArray{TO,DO}, μ::AbstractArray{TP,DP}, logσ2::AbstractArray{TO}; + i_finobs = .! isnan.(obs), # TODO provide with y_unc to save allocations + σfac=one(TO)) where {TO, TP, DO, DP} + sum( + i_finobs .* (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) + ) +end + + function neg_logden_indep_normal(obs::AbstractGPUArray, μ::AbstractGPUArray, logσ2::AbstractGPUArray{ET}; σfac=one(ET)) where ET error("update neg_logden_indep_normal for GPU to each solsum rather than total sum") diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 0df7101..45e0aca 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -15,6 +15,7 @@ using Bijectors #import PreallocationTools as PAT import JLD2 import PDMats +import StaticArrays as SA import Lux # import Enzyme @@ -36,7 +37,7 @@ n_θM = length(par_templates.θM) n_cov = get_hybridproblem_n_covar(prob; scenario) pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) n_covP = length(pbm_covars) - +pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_templates.θP)) for k in pbm_covars]) g, ϕg = get_hybridproblem_MLapplicator(prob; scenario) @@ -47,15 +48,19 @@ n_covP = length(pbm_covars) transP, transM = get_hybridproblem_transforms(prob; scenario) # cor_ends = get_hybridproblem_cor_ends(prob; scenario) + approx1 = get_hybridproblem_HVIApproximation(prob; scenario) + # TODO update init_hybrid_ϕunc + ϕq, approx = init_hybrid_ϕunc(approx1, cor_ends; transM, par_templates.θM, n_site) ϕqPc1 = CA.ComponentVector( μζP = log.(CA.getdata(par_templates.θP)), - logσ_ζP = ones(n_θP) .* log(0.01), - ρsP = [0.1], + #logσ_ζP = ones(n_θP) .* log(0.01), + logσ_ζP = ϕq.logσ_ζP, + ρsP = ϕq.ρsP, ) ϕqIc1 = CA.ComponentVector( - logσ_ζM = ones(n_θM) .* log(0.01), - ρsM = [0.1], + logσ_ζM_offsets = ϕq.logσ_ζM_offsets, + ρsM = ϕq.ρsM, ) ϕqP = ϕqP2 = CA.getdata(ϕqPc1) ϕqI = CA.getdata(ϕqIc1) @@ -64,7 +69,7 @@ n_covP = length(pbm_covars) # indiv_args1 = (;xM = CA.getdata(xM)[:,i_indiv_batch], xP = CA.getdata(xP)[:,i_indiv_batch], y_o = CA.getdata(y_o)[:,i_indiv_batch], y_unc = CA.getdata(y_unc)[:,i_indiv_batch], - indiv_train = i_indiv_batch) + i_indiv_train = i_indiv_batch) # pbm = create_nsite_applicator(get_hybridproblem_PBmodel(prob; scenario), n_MC; diffchunk) nljoint_args_fix =(; @@ -72,13 +77,14 @@ n_covP = length(pbm_covars) fLogLik = get_hybridproblem_neg_logden_obs(prob; scenario) ) - ϕqIcS = CA.ComponentVector(;logσ_ζM_offsets = ϕqIc1.logσ_ζM[2:end] .- ϕqIc1.logσ_ζM[1], ϕqIc1.ρsM) - intϕqIS = get_concrete(ComponentArrayInterpreter(ϕqIcS)) - approxS_ = CP.MeanUniScalingHVIApproximation(ϕqIc1.logσ_ζM[1]) - approxS = get_hybridproblem_HVIApproximation(prob; scenario) - n_M = get_numberof_MLinputs(approxS, par_templates.θM) - h0S = CP.prepare_elbo_helpers(approxS, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); - randn!(h0S.θsP) # for testing should initialized to finite values + # ϕqIcS = CA.ComponentVector(;logσ_ζM_offsets = ϕqIc1.logσ_ζM[2:end] .- ϕqIc1.logσ_ζM[1], ϕqIc1.ρsM) + # intϕqIS = get_concrete(ComponentArrayInterpreter(ϕqIcS)) + # approxS_ = CP.MeanUniScalingHVIApproximation(ϕqIc1.logσ_ζM[1]) + # approxS = get_hybridproblem_HVIApproximation(prob; scenario) + n_M = get_numberof_MLinputs(approx, par_templates.θM) + h0 = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC, n_cov, n_covP, n_M, cor_ends, n_obs); + randn!(h0.θsP) # for testing should initialized to finite values + rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) #------------------------------------------------------------- @@ -183,359 +189,10 @@ n_covP = length(pbm_covars) ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) end -@testset "pullback_g_apply!" begin - # @test ϕms0 == ϕms0z - # @test ϕms2 == ϕms2z - # @test size(ϕms0) == (n_M, n_site) - # @test size(ϕms2) == (n_M, n_MC, n_site) - # () -> begin # gradient(sum) - # gr_zygote = Zygote.gradient((ϕg2) -> sum(CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP)), ϕg2v ) - # s, pullback_s_zygote = Zygote.pullback((ϕg2) -> sum(CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP)), ϕg2v ) - # #gr_zygote2 = pullback_s_zygote(ones(eltype(ϕg2v), size(ϕg2v)...)) - # gr_zygote2 = pullback_s_zygote(one(eltype(ϕg2v))) - # @test gr_zygote2[1] ≈ gr_zygote[1] - # y, pullback_zygote = Zygote.pullback((ϕg2) -> CP.g_apply_oop(ϕg2, xM, ζP, pbm_covar_indices2, g2, xMP), ϕg2v ) - # gr_zygote3 = pullback_zygote(ones(eltype(y), size(y)...)) - # gr_zygote3[1] ≈ gr_zygote[1] - # end - # # - # # concatenate function f3(g(ϕ_g)) - # f3 = (x) -> sum(3.0 .* x) - # CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP) - # # one pass of composed function - # gr_zygote = Zygote.gradient((ϕg2, ζsP) -> f3(CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP)), ϕg2v, ζsP ) - # s, pullback_s_zygote = Zygote.pullback((ϕg2) -> f3(CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP)), ϕg2v ) - # gr_zygote2 = pullback_s_zygote(one(eltype(s))) - # @test gr_zygote2[1] ≈ gr_zygote[1] - # # mixed AD, differentiate f3 by FowardDiff and pull back through g - # y_oop = CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP) - # #gr_h = Zygote.gradient(y -> sum(f3(y)), y1)[1] - # gr_h = ForwardDiff.gradient(y -> sum(f3(y)), y_oop) - # y, pullback_zygote = Zygote.pullback((ϕg2) -> CP.g_apply_oop(ϕg2, xM, ζsP, pbm_covar_indices2, g2, xMP), ϕg2v ) - # @test y == y_oop - # gr_zygote3 = pullback_zygote(gr_h) - # gr_zygote3[1] ≈ gr_zygote[1] - # # - # dϕg = zero(ϕg2v) - # dζsP = zero(ζsP) - # #Enzyme.make_one!(y) - # y .= rand() - # dϕg .= rand() # check that initial values do not effect result - # dζsP .= rand() # check that initial values do not effect result - # dy = convert.(eltype(y), gr_h) - # CP.pullback_g_apply!(y, dϕg, dζsP, dy, ϕg2v, xM, ζsP, pbm_covar_indices2, g2, h) - # @test y == y_oop - # @test dϕg ≈ gr_zygote[1] - # @test dζsP ≈ gr_zygote[2] rtol=1e-3 - # @test dy ≈ gr_h # not modified - # #@benchmark CP.pullback_g_apply!(y, dϕg, dy, ϕg2v, xM, ζsP, pbm_covar_indices2, g2, h) - # # - # () -> begin # explicitly splitting the forward and backward pass - # # they get cached anymay and require allocating the Duplicated Wrappers twice - # # hence there is no performance benefit - # # Compile once outside the hot loop - # fwd, rev = Enzyme.autodiff_thunk( - # Enzyme.ReverseSplitNoPrimal, - # Enzyme.Const{typeof(g_apply!)}, - # Enzyme.Const, - # Enzyme.Duplicated{typeof(y)}, - # Enzyme.Duplicated{typeof(ϕg2v)}, - # Enzyme.Const{typeof(xM)}, - # Enzyme.Const{typeof(ζP)}, - # Enzyme.Const{typeof(pbm_covar_indices2)}, - # Enzyme.Const{typeof(g2)}, - # Enzyme.Duplicated{typeof(h.xMP)} - # ) - # # take care, dy is also modified - # function grad2_g_apply!(y, dϕg, dy, ϕg, xM, ζP, pbm_covar_indices, g, h, fwd, rev) - # fill!(dϕg, zero(eltype(dϕg))) - # fill!(h.dxMP, zero(eltype(h.dxMP))) - # copyto!(h.dy, dy) # copy to avoid modifying dy - # tape, _, _ = fwd( - # Enzyme.Const(g_apply!), - # Enzyme.Duplicated(y, h.dy), - # Enzyme.Duplicated(ϕg, dϕg), - # Enzyme.Const(xM), - # Enzyme.Const(ζP), - # Enzyme.Const(pbm_covar_indices), - # Enzyme.Const(g), - # Enzyme.Duplicated(h.xMP, h.dxMP) - # ) - # rev( - # Enzyme.Const(g_apply!), - # Enzyme.Duplicated(y, h.dy), - # Enzyme.Duplicated(ϕg, dϕg), - # Enzyme.Const(xM), - # Enzyme.Const(ζP), - # Enzyme.Const(pbm_covar_indices), - # Enzyme.Const(g), - # Enzyme.Duplicated(h.xMP, h.dxMP), - # tape - # ) - # return nothing - # end - - # dy = convert.(eltype(y), gr_h) - # grad2_g_apply!(y, dϕg, dy, ϕg2v, xM, ζP, pbm_covar_indices2, g2, h, fwd, rev) - # @test y == y_oop - # @test dϕg ≈ gr_zygote[1] - # @test dy ≈ gr_h # not modified - # #@usingany BenchmarkTools - # #@benchmark grad2_g_apply!(y, dϕg, dy, ϕg2v, xM, ζP, pbm_covar_indices2, g2, h, fwd, rev) - # end -end - -@testset "pullback_sample_ζsP!" begin - # ϕqPc = intϕqP(ϕqP) - # rnormP = zero(ζsP) - # randn!(rnormP) # before input gaussian noise - # ζsP .= 0 - # #logσ_ζP = zero(ϕqPc.logσ_ζP) # cretes a view rather than copy - # logσ_ζP = zero(ϕqPc.logσ_ζP) - # CP.sample_ζsP!(ζsP, logσ_ζP, rnormP, ϕqPc) - # mean(ζsP; dims=2) - # ζsP1 = copy(ζsP) - - # # Enzyme result via the mutating routine (2-D: n_θP * n_MC × n_in) - # dζsP = zero(ζsP) .+ one(eltype(ζsP)) - # dlogσ_ζP = zero(logσ_ζP) .+ one(eltype(ζsP)) - # dϕqc = zero(ϕqPc) - - # randn!(dϕqc) # test that is zerod inside pullback - # dζsP_ = copy(dζsP) - # rnormP_ = copy(rnormP) - # logσ_ζP_ = copy(logσ_ζP) - # dlogσ_ζP_ = copy(dlogσ_ζP) - # #CP.pullback_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, rnormP, logσ_ζP, ϕqPc) # needs rnormP to be noise - # CP.pullback_sample_ζsP!(dϕqc, dζsP, dlogσ_ζP, ζsP, logσ_ζP, rnormP, ϕqPc) - # @test ζsP == ζsP1 # same forward result - # @test rnormP == rnormP_ - # @test dζsP == dζsP_ - # @test dlogσ_ζP == dlogσ_ζP_ - # @test logσ_ζP == logσ_ζP_ - # # without correlation - # #@test all(dϕqc[Val(:μζP)] .== n_MC) - # # #@test dϕqc[Val(:logσ_ζP)] ≈ vec(sum(rnormP; dims=2)) # - # dϕqc_comb = copy(dϕqc) - - # randn!(ζsP) # test initial not relevant - # pb_sample_ζsP = CP.primal_pullback_sample_ζsP!(ζsP, logσ_ζP, rnormP, ϕqPc) - # @test ζsP == ζsP1 # same forward result - # @test rnormP ≈ rnormP_# computed the forward pass - # @test logσ_ζP == logσ_ζP_ - # dϕqc .= 0.1 # test initial value not relevant - # #dϕqc .= 0.01 # should not influence results - # pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) - # #pb_sample_ζsP(rnormP, logσ_ζP) - # @test rnormP ≈ rnormP_ # did not modify - # @test logσ_ζP == logσ_ζP_ # not modified - # @test dlogσ_ζP == dlogσ_ζP_ # not modified - # @test dζsP == dζsP_ - # #hcat(dϕqc, dϕqc_comb) - # @test CA.getdata(dϕqc) ≈ CA.getdata(dϕqc_comb) - # # - # # test another pullback - # #dζsP .= dζsP * eltype(dζsP)(2) - # pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) - # @test CA.getdata(dϕqc) ≈ CA.getdata(dϕqc_comb) - # # - # # @usingany BenchmarkTools - # # @benchmark pb_sample_ζsP(dϕqc, dζsP, dlogσ_ζP) -end - -# to be differentiated, return only first entry of result -function _ftmp2(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) - CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] -end - -function grad_neg_elbo_sites_enzyme() # differentiate entire neg_elbo_sites by enzyme - # do not use DiffCache here for helpers_sites - h0p = CP.prepare_elbo_helpers(approx, ϕg, ϕqP; - n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP0, n_M, cor_ends, n_obs, - use_diff_cache = Val(false)) - # and store results to compare to hand-crafted mixed AD - dh0p = Enzyme.make_zero(h0p) - @test dh0p !== h0p # real copy rather than reference - # - indiv_args = indiv_args1 - ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) - dϕ = map(zero, ϕ) - sample_args = (; - approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, - transP, transM) - # - rng1 = StableRNG(1234) - CP.randnPM!(rng1, rnormPM) - randn!(rng1, ϕ.ϕg) - randn!(rng1, indiv_args.xM) - primal_enz = tmp = _ftmp2(h0p, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) - Enzyme.make_zero!(dϕ) - Enzyme.make_zero!(dh0p) - Enzyme.autodiff( - Enzyme.set_runtime_activity(Enzyme.Reverse) , - _ftmp2, - Enzyme.Active, - Enzyme.Duplicated(h0p, dh0p), - Enzyme.Duplicated(ϕ, dϕ), - Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), - Enzyme.Const(sample_args), - Enzyme.Const(indiv_args), - Enzyme.Const(nljoint_args_fix), - ) - () -> begin - dϕg0_enz ≈ dϕ.ϕg - dϕqP0_enz ≈ dϕ.ϕqP - dϕqI0_enz ≈ dϕ.ϕqI - hcat(dϕg0_enz, dϕ.ϕg) - end - dϕg0_enz = copy(dϕ.ϕg) - dϕqP0_enz = copy(dϕ.ϕqP) - dϕqI0_enz = copy(dϕ.ϕqI) - () -> begin - #@usingany JLD2 - #fname = "intermediate/test_enzyme_dphi0.jld2" - fname = "intermediate/test_enzymeT_dphi0.jld2" - mkpath("intermediate") - JLD2.jldsave(fname, false, IOStream; primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz) - primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz = JLD2.load(fname, - "primal_enz", "dϕg0_enz", "dϕqP0_enz", "dϕqI0_enz"); - end - - h2p = CP.prepare_elbo_helpers(approx, ϕg2, ϕqP; - n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, n_obs, - use_diff_cache = Val(false)) - dh2p = Enzyme.make_zero(h2p) - indiv_args = indiv_args1 - ϕ2 = (;ϕg=ϕg2v, ϕqP, ϕqI) - dϕ2 = map(zero, ϕ2) - sample_args2 = (; - approx, g=g2, is_testmode, pbm_covar_indices = pbm_covar_indices2, intϕqP, intϕqI, cor_ends, - transP, transM) - # - rng1 = StableRNG(1234) - CP.randnPM!(rng1, rnormPM) - randn!(rng1, ϕ2.ϕg) - randn!(rng1, indiv_args.xM) - primal_enz = _ftmp2(h2p, ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix) - Enzyme.make_zero!(dϕ2) - Enzyme.make_zero!(dh2p) - Enzyme.autodiff( - Enzyme.set_runtime_activity(Enzyme.Reverse) , - _ftmp2, - Enzyme.Active, - Enzyme.Duplicated(h2p, dh2p), - Enzyme.Duplicated(ϕ2, dϕ2), - Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), - Enzyme.Const(sample_args2), - Enzyme.Const(indiv_args), - Enzyme.Const(nljoint_args_fix), - ) - () -> begin - primal2_enz ≈ primal_enz - dϕg2_enz ≈ dϕ2.ϕg - @test dϕqP2_enz ≈ dϕ2.ϕqP atol=0.001 - dϕqI2_enz ≈ dϕ2.ϕqI - hcat(dϕg2_enz, dϕ2.ϕg) - hcat(dϕqP2_enz, dϕ2.ϕqP) - end - dϕg2_enz = copy(dϕ2.ϕg) - dϕqP2_enz = copy(dϕ2.ϕqP) - dϕqI2_enz = copy(dϕ2.ϕqI) - () -> begin - #fname = "intermediate/test_enzyme_dphi2.jld2" - fname = "intermediate/test_enzymeT_dphi2.jld2" - mkpath("intermediate") - JLD2.jldsave(fname, false, IOStream; primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz) - primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz = JLD2.load(fname, - "primal2_enz", "dϕg2_enz", "dϕqP2_enz", "dϕqI2_enz"); - end -end - -function grad_neg_elbo_sites_enzyme_scaled() - # now with approxS and respective arguments - is_testmode = false - indiv_args = indiv_args1 - ϕ2S = (;ϕg=ϕg2v, ϕqP, ϕqI=CA.getdata(ϕqIcS)) - dϕ2S = map(zero, ϕ2S) - sample_args2S = (; - approx=approxS, g=g2, is_testmode, pbm_covar_indices = pbm_covar_indices2, - intϕqP, intϕqI=intϕqIS, cor_ends, - transP, transM) - h2Sp = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; - n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, use_diff_cache = Val(false)) - dh2Sp = Enzyme.make_zero(h2Sp) - # - rng1 = StableRNG(1234) - CP.randnPM!(rng1, rnormPM) - randn!(rng1, ϕ2S.ϕg) - randn!(rng1, indiv_args.xM) - primal_enz2S = tmp = _ftmp2(h2Sp, ϕ2S, rnormPM, sample_args2S, indiv_args) - Enzyme.make_zero!(dϕ2S) - Enzyme.make_zero!(dh2Sp) - Enzyme.autodiff( - Enzyme.set_runtime_activity(Enzyme.Reverse) , - _ftmp2, - Enzyme.Active, - Enzyme.Duplicated(h2Sp, dh2Sp), - Enzyme.Duplicated(ϕ2S, dϕ2S), - Enzyme.DuplicatedNoNeed(rnormPM, Enzyme.make_zero(rnormPM)), - Enzyme.Const(sample_args2S), - Enzyme.Const(indiv_args), - ) - dϕg2S_enz = copy(dϕ2S.ϕg) - dϕqP2S_enz = copy(dϕ2S.ϕqP) - dϕqI2S_enz = copy(dϕ2S.ϕqI) - () -> begin - #fname = "intermediate/test_enzyme_dphi2.jld2" - fname = "intermediate/test_enzymeS_dphi2.jld2" - mkpath("intermediate") - JLD2.jldsave(fname, false, IOStream; primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz) - primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load(fname, - "primal_enz2S", "dϕg2S_enz", "dϕqP2S_enz", "dϕqI2S_enz"); - end -end - -function tmp_reference_revision() - #@usingany ReferenceRevision - ref = open_process(rev="origin/sitepar", path="tmp/ref", use=[:Revise, :ShareAdd]) - ref.eval(:(Pkg.activate("tmp/ref/test"))) - #ref.eval(:(Pkg.instantiate())) - #ref.Pkg.project().path - ref.eval(:(using HybridVariationalInference)) - ref.eval(:(import HybridVariationalInference as HVI)) - ref.eval(:(using Lux)) # to deserialize objects - ref.eval(:(using SharedArrays)) # to deserialize objects - #ref.eval(:(@usingany Infiltrator)) # not output during infiltrat in revision - tmp = ref.HVI.test3(3) - ref.Revise.retry() # after every change of source file - res0v, gradh0v = ref.HVI.grad_neg_elbo_sites( - h0, - (;), - approx, - rnormPM, - ϕgv, ϕqP, ϕqI, g, nothing; #pbm_covar_indices2; - i_sites_train = 1:n_site, - intϕqP, intϕqI, - xM, - cor_ends, - transP, transM, - is_testmode = false, - ) - close(ref) - - map(k -> k => res0v[k] ≈ res0[k], keys(res0)) - hcat(res0v.dϕqIc, res0.dϕqIc) -end - @testset "grad_neg_elbo_sites" begin ϕqIc = intϕqI(ϕqI) ϕqPc = intϕqP(ϕqP) is_testmode = false - n_threads_proc = min(Threads.nthreads(), 4) - n_workers = Distributed.nworkers() * n_threads_proc - basesize = n_site ÷ Distributed.nworkers() - distributedEx = Transducers.DistributedEx(;basesize, - threads_basesize = Int(ceil(basesize / n_threads_proc))) # rng1 = StableRNG(1234) CP.randnPM!(rng1, rnormPM) @@ -544,9 +201,10 @@ end indiv_args = indiv_args1 ϕ = (;ϕg=ϕgv, ϕqP, ϕqI) sample_args = (; - approx, g, is_testmode, pbm_covar_indices = nothing, intϕqP, intϕqI, cor_ends, + approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) primal = CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + res0, gradh0 = CP.grad_neg_elbo_sites( #@descend_code_warntype CP.neg_elbo_sites!( h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) From bc395a49ab9c45df3bc7c7f0cf9c5debe667af46 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Tue, 20 Oct 2026 14:23:39 +0000 Subject: [PATCH 3/4] make NormalScalingModelApplicator work with Enzyme needed to replace broadcast by explicit loop --- src/DoubleMM/f_doubleMM.jl | 4 +- src/ModelApplicator.jl | 31 +++- src/PBMApplicator.jl | 8 +- src/elbo_site.jl | 8 +- src/elbo_site_grad.jl | 6 +- src/logden_normal.jl | 10 +- test/test_elbo_site_doubleMM.jl | 304 +++++++++++++------------------- 7 files changed, 166 insertions(+), 205 deletions(-) diff --git a/src/DoubleMM/f_doubleMM.jl b/src/DoubleMM/f_doubleMM.jl index 0b073fd..afd8b7f 100644 --- a/src/DoubleMM/f_doubleMM.jl +++ b/src/DoubleMM/f_doubleMM.jl @@ -155,7 +155,7 @@ end function f_doubleMM_pop!(pred, θs, xPc::CA.ComponentVector{T}) where T S1 = view(xPc, Val(:S1), :) S2 = view(xPc, Val(:S2), :) - is_valid = isfinite.(S1) .&& isfinite.(S2) + # is_valid = isfinite.(S1) .&& isfinite.(S2) @row_valviews r0, r1, K1, K2 = θs # each variable is a matrix (n_obs x n_site) pred .= r0' .+ r1' .* S1 ./ (K1' .+ S1) .* S2 ./ (K2' .+ S2) @@ -243,7 +243,7 @@ function HVI.get_hybridproblem_MLapplicator( # TODO make this dependent on approx 1:length(θM) else - 1:0 + nothing end g = if (:use_rangescaling ∈ scen) RangeScalingModelApplicator(g_nomag, lowers, uppers, eltype(ϕ_g0); range_scaled) diff --git a/src/ModelApplicator.jl b/src/ModelApplicator.jl index 34bf5c7..0e24781 100644 --- a/src/ModelApplicator.jl +++ b/src/ModelApplicator.jl @@ -177,7 +177,7 @@ It usually corresponds to the type used in other ML-parts of the model, e.g. `Fl """ function NormalScalingModelApplicator( app::AbstractModelApplicator, lowers, uppers, FT::Type; - range_scaled = 1:0, + range_scaled = nothing, #1:0, repeat_inner::Integer = 1) pars = map(lowers, uppers) do lower, upper dζ = fit(Normal, @qp_l(lower), @qp_u(upper)) @@ -217,7 +217,7 @@ function apply_model(app::NormalScalingModelApplicator, x, ϕ; kwargs...) # @show typeof(app.μ) # @show typeof(ϕ) @assert eltype(app.μ) == eltype(ϕ) - ans = if !isempty(app.range_scaled) + ans = if !isnothing(app.range_scaled) ans_scaled = norminvcdf.(app.μ, app.σ, index_firstdim(y_perc,app.range_scaled)) # from StatsFuns combine_range(y_perc, ans_scaled, app.range_scaled) else @@ -236,8 +236,23 @@ function apply_model!(y, app::NormalScalingModelApplicator, x, ϕ; kwargs...) # @show typeof(app.μ) # @show typeof(ϕ) @assert eltype(app.μ) == eltype(ϕ) - if !isempty(app.range_scaled) - view_firstdim(y,app.range_scaled) .= norminvcdf.(app.μ, app.σ, view_firstdim(y,app.range_scaled)) # from StatsFuns + if !isnothing(app.range_scaled) + for x in app.range_scaled; @assert(x ∈ axes(y,1)); end + # .= fails with Enzyme + # The destination is dotview(y, r, :), a view aliasing the active array y. + # The in-place materialize! path runs preprocess → broadcast_unalias → unalias → + # mightalias. Because the destination view may alias an RHS operand (view(y,r,:)), + # Julia inserts a runtime aliasing branch: if they alias, one operand is copied into + # a fresh temporary Vector{Float32}. That conditionally-allocated temp is + # active-or-constant depending on runtime flow — precisely the "Mismatched activity" + # phi node and "Unknown object of type Vector{Float32}" + #y[app.range_scaled,:] .= norminvcdf.(app.μ, app.σ, view(y,app.range_scaled,:)) + # = works with Enzyme but allocates a fresh array + #y[app.range_scaled,:] = norminvcdf.(app.μ, app.σ, view(y,app.range_scaled,:)) + # need to code the loop by hand + @inbounds for j in axes(y,2), (i,ri) in enumerate(app.range_scaled) + y[ri,j] = norminvcdf(app.μ[i], app.σ[i], y[ri,j]) + end else y .= norminvcdf.(app.μ, app.σ, y) end @@ -245,10 +260,10 @@ function apply_model!(y, app::NormalScalingModelApplicator, x, ϕ; kwargs...) end -index_firstdim(v::AbstractVector, i) = v[i] -index_firstdim(v::AbstractMatrix, i) = v[i,:] -view_firstdim(v::AbstractVector, i) = view(v, i) -view_firstdim(v::AbstractMatrix, i) = view(v,i,:) +@inline index_firstdim(v::AbstractVector, i) = v[i] +@inline index_firstdim(v::AbstractMatrix, i) = v[i,:] +@inline view_firstdim(v::AbstractVector, i) = view(v, i) +@inline view_firstdim(v::AbstractMatrix, i) = view(v,i,:) """ RangeScalingModelApplicator(app, y0) diff --git a/src/PBMApplicator.jl b/src/PBMApplicator.jl index dfb8932..e8b1ab4 100644 --- a/src/PBMApplicator.jl +++ b/src/PBMApplicator.jl @@ -91,24 +91,24 @@ function apply_model!(pred, app::PBMPopulationApplicator{AXS,AXP}, θsP::Abstrac pred end -function (app::PBMPopulationApplicator)(θP::AbstractMatrix, θMs_tr::AbstractArray, xP::AbstractMatrix) +function (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) # @assert n_MC = size(θP,2) - pred = zeros(n_obs, n_site, n_MC) + pred = zeros(FT, n_obs, n_site, n_MC) for i_site in 1:n_site HVI.apply_model!(pred[:,i_site,:], app, θP, θMs_tr[i_site,:,:], xP[:,i_site]) end (pred, nothing) end -function (app::PBMPopulationApplicator)(θP::AbstractVector, θMs_tr::AbstractMatrix, xP::AbstractMatrix) +function (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]) n_site, n_θM = size(θMs_tr) θsP = reshape(θP, length(θP), 1) - pred = zeros(n_obs, n_site, 1) + pred = zeros(FT, n_obs, n_site, 1) # i_site = 1 for i_site in 1:n_site θsM = reshape(view(θMs_tr,i_site,:), n_θM, 1) diff --git a/src/elbo_site.jl b/src/elbo_site.jl index f933cec..24c7a5a 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -44,10 +44,11 @@ function neg_elbo_sites!( ϕm_it = eachslice(h[ϕms_buffer_key]; dims = ndims(h[ϕms_buffer_key])) template = ϕqI # only important for gradient θsP = h.θsP - keys_nljoint_args_indi = keys(nljoint_args_inds) + keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) # closure with approx, transM, and kwargs function compute_elboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) - nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup) + #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) compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm, ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) end @@ -61,9 +62,10 @@ function neg_elbo_sites!( # costTrans = sum(x -> x.costTrans, res_site) #elbo = sum(first, res_site) - sum(h.logσ_ζP) elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) - (; elbo, ζsP=copy(h.ζsP), ϕm=copy(h[ϕms_buffer_key])) end +add_names(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) + diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 8c45667..29c80d2 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -43,8 +43,8 @@ function grad_neg_elbo_sites( nljoint_args_indi = NamedTuple{keys(nljoint_args_inds)}( first(zip_eachlastdims(nljoint_args_inds))) gradh = !isempty(grad_elbo_helpers) ? grad_elbo_helpers : - Test.@inferred prepare_gradelbo_helpers( - #prepare_gradelbo_helpers( + #Test.@inferred prepare_gradelbo_helpers( + prepare_gradelbo_helpers( ϕqIc, selectdim(h_ϕm, ndims(h_ϕm), 1), h.ζsP, ϕg, ϕqPc, approx; pbm_covar_indices, n_workers, h, rnormMi = rnormPM.M[1], nljoint_args_indi, @@ -349,7 +349,7 @@ function get_pullback_g_apply(::AbstractArray{TG}, ::AbstractArray{TF}; # assert that buffers were constructed with correct sizes n_cov_f, n_site_f = size(xM) n_covP_f = isnothing(pbm_covar_indices) ? 0 : length(pbm_covar_indices) - n_MC_f = size(ζsP,1) + n_MC_f = size(ζsP,2) @assert (n_cov, n_covP, n_MC, n_site) == (n_cov_f, n_covP_f, n_MC_f, n_site_f) @assert size(dϕms_buffer) == size(dϕm) diff --git a/src/logden_normal.jl b/src/logden_normal.jl index e9b04ff..679bdfc 100644 --- a/src/logden_normal.jl +++ b/src/logden_normal.jl @@ -67,11 +67,13 @@ function neg_logden_indep_normal(obs::AbstractArray{TO,DO}, μ::AbstractArray{TP end -function neg_logden_indep_normal_sum(obs::AbstractArray{TO,DO}, μ::AbstractArray{TP,DP}, logσ2::AbstractArray{TO}; - i_finobs = .! isnan.(obs), # TODO provide with y_unc to save allocations - σfac=one(TO)) where {TO, TP, DO, DP} +function neg_logden_indep_normal_sum(obs::AbstractArray{TO,DO}, μ::AbstractArray, logσ2::AbstractArray; + #i_finobs = .! isnan.(obs), # TODO provide with y_unc to save allocations + σfac=one(TO)) where {TO, DO} + return sum(obs) sum( - i_finobs .* (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) + #i_finobs .* (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) + (σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) ./ eltype(μ)(2) ) end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 45e0aca..1ba949b 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -88,111 +88,111 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te #------------------------------------------------------------- -@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 "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) - is_testmode = false + is_testmode = true # rng1 = StableRNG(1234) CP.randnPM!(rng1, rnormPM) @@ -204,97 +204,39 @@ end approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) primal = CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) - + function loop_neg_elbo_sites!(n, h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + for _ in 1:n + @noinline CP.neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + end + nothing + end + #@profview_allocs loop_neg_elbo_sites!(100_000, h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + function alloc_neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + @allocated loop_neg_elbo_sites!(10, h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + end + # allocations come from Lux StatefulLuxLayer + alloc_neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + @test alloc_neg_elbo_sites!(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) <= 128_000 + # + sample_args_rev = (; sample_args..., is_testmode=false) res0, gradh0 = CP.grad_neg_elbo_sites( #@descend_code_warntype CP.neg_elbo_sites!( h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) - res0_, gradh0_ = CP.grad_neg_elbo_sites( # test deterministic result and distributed - h0, (;), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val, - executor = distributedEx - ) - @test all(map(≈, res0_, res0)) - # # if we saved Enzyme results earlier, compare to them - if isfile("intermediate/test_enzymeT_dphi0.jld2") - primal_enz, dϕg0_enz, dϕqP0_enz, dϕqI0_enz = JLD2.load( - "intermediate/test_enzymeT_dphi0.jld2", - "primal_enz", "dϕg0_enz", "dϕqP0_enz", "dϕqI0_enz"); - @test primal_enz ≈ primal[1] - @test dϕg0_enz ≈ res0.dϕg - @test dϕqI0_enz ≈ res0.dϕqIc - @test dϕqP0_enz ≈ res0.dϕqPc - #hcat(dϕqI0_enz, CA.getdata(res0.dϕqIc)) - #hcat(dϕqP0_enz, CA.getdata(res0.dϕqPc)) - #hcat(dϕg0_enz, CA.getdata(res0.dϕg)) - #dϕqP0_enz - CA.getdata(res0.dϕqP) - end - # - #---------------- matrix mode with population covariates - ϕ2 = (;ϕ..., ϕg=ϕg2v) - sample_args2 = (;sample_args..., g=g2, pbm_covar_indices = pbm_covar_indices2) - rng1 = StableRNG(1234) - CP.randnPM!(rng1, rnormPM) - randn!(rng1, ϕ2.ϕg) - randn!(rng1, indiv_args.xM) - primal2 = CP.neg_elbo_sites!(h2, ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix)[1] - res0, gradh2 = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val) - res0_, gradh2_ = CP.grad_neg_elbo_sites(h2, (;), ϕ2, rnormPM, sample_args2, indiv_args, nljoint_args_fix; n_MC_val, - executor = distributedEx,) - @test all(map(≈, res0_, res0)) - # if we saved Enzyme results earlier, compare to them - if isfile("intermediate/test_enzymeT_dphi2.jld2") - primal2_enz, dϕg2_enz, dϕqP2_enz, dϕqI2_enz = JLD2.load( - "intermediate/test_enzymeT_dphi2.jld2", - "primal2_enz", "dϕg2_enz", "dϕqP2_enz", "dϕqI2_enz"); - @test primal2_enz ≈ primal2[1] - @test dϕg2_enz ≈ res0.dϕg - @test dϕqI2_enz ≈ res0.dϕqIc - @test dϕqP2_enz ≈ res0.dϕqPc atol = 0.001 - #hcat(dϕqP2_enz, res0.dϕqP) - end - #---------------- approxM with non-empty h.sample_buffers and scaling - # and using eachcol in PBM function - ϕ2S = (;ϕ2..., ϕqI=CA.getdata(ϕqIcS)) - sample_args2S = (;sample_args2..., approx=approxS, intϕqI=intϕqIS) - #nljoint_args_fix2S = nljoint_args_fix =(; f! = pbm_col) - nljoint_args_fix2S = (; nljoint_args_fix..., f! = pbm_col) - # h2Sp = CP.prepare_elbo_helpers(sample_args2S.approx, ϕg2, ϕqP; - # n_θP, n_θM, n_site, n_MC, n_cov, n_covP = n_covP2, n_M, cor_ends, use_diff_cache = Val(false)) - rng1 = StableRNG(1234) - CP.randnPM!(rng1, rnormPM) - randn!(rng1, ϕ2S.ϕg) - randn!(rng1, indiv_args.xM) - primal2S = CP.neg_elbo_sites!(h2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) - res0, gradh2S = CP.grad_neg_elbo_sites( - h2S, (;), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, - #executor = distributedEx, - ) - if isfile("intermediate/test_enzymeS_dphi2.jld2") - primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load( - "intermediate/test_enzymeS_dphi2.jld2", - "primal_enz2S", "dϕg2S_enz", "dϕqP2S_enz", "dϕqI2S_enz"); - @test primal_enz2S ≈ primal2S[1] - @test dϕg2S_enz ≈ res0.dϕg - @test dϕqI2S_enz ≈ res0.dϕqIc atol=0.01 - @test dϕqP2S_enz ≈ res0.dϕqPc atol=0.01 - #hcat(dϕqP2_enz, res0.dϕqP) - #hcat(dϕqI2S_enz, res0.dϕqI) - end - function loop_grad_neg_elbo_sites( - n, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + n, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) for _ in 1:n @noinline CP.grad_neg_elbo_sites( - h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix; n_MC_val) + h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) end nothing end - function alloc_grad_neg_elbo_sites(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) + 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, h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) - #@profview_allocs loop_grad_neg_elbo_sites(1_000,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix) - @test (@allocated loop_grad_neg_elbo_sites(100,h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix)) <= 3_264_000 #3_260_800 + loop_grad_neg_elbo_sites(1, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + #@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(h2S, gradh2S, ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) + 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) From eed2e22527ba9cc58ec481fb01079d9c6110ff51 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Tue, 20 Oct 2026 15:12:18 +0000 Subject: [PATCH 4/4] rename compute_elboi_z_cl! to compute_nelboi_z_cl! --- src/ModelApplicator.jl | 6 +++++- src/elbo_site.jl | 6 +++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/ModelApplicator.jl b/src/ModelApplicator.jl index 0e24781..6a9a95e 100644 --- a/src/ModelApplicator.jl +++ b/src/ModelApplicator.jl @@ -249,10 +249,14 @@ function apply_model!(y, app::NormalScalingModelApplicator, x, ϕ; kwargs...) #y[app.range_scaled,:] .= norminvcdf.(app.μ, app.σ, view(y,app.range_scaled,:)) # = works with Enzyme but allocates a fresh array #y[app.range_scaled,:] = norminvcdf.(app.μ, app.σ, view(y,app.range_scaled,:)) - # need to code the loop by hand + # need to code the loop by hand, there activity is clear @inbounds for j in axes(y,2), (i,ri) in enumerate(app.range_scaled) y[ri,j] = norminvcdf(app.μ[i], app.σ[i], y[ri,j]) end + # shortcircuits the might_alias but Enzyme still gets a mismatched activity of + # the not-taken runtime branch (activity differs on wheter its copied or not) + # v = view(y,app.range_scaled,:) + # v .= norminvcdf.(app.μ, app.σ, v) else y .= norminvcdf.(app.μ, app.σ, y) end diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 24c7a5a..a9a07be 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -46,7 +46,7 @@ function neg_elbo_sites!( θsP = h.θsP keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) # closure with approx, transM, and kwargs - function compute_elboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) + 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) compute_nelboi_z!(hi, approx, rnormM, nljoint_args_indi, ϕm, @@ -56,7 +56,7 @@ function neg_elbo_sites!( #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 - elbo_z = mapreduce(compute_elboi_z_cl!, +, h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it) + elbo_z = mapreduce(compute_nelboi_z_cl!, +, h.helpers_sites, rnormPM.M, zip_eachlastdims(nljoint_args_inds), ϕm_it) # E = sum(x -> x.E, res_site) # loglik = sum(x -> x.loglik, res_site) # costTrans = sum(x -> x.costTrans, res_site) @@ -85,7 +85,7 @@ function compute_nelboi_z!(hi, approx::AbstractHVIApproximation, exp_ladJacTM = transformζ!(hi.θsM, transM, hi.ζsM) # return value captures ladJacT # first component needs to be the full elbo exp_nL = exp_nLi!(hi.y, θsP, hi.θsM, nljoint_args_indi, nljoint_args_fix)[1] - elbozi = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM) + nelboi_z = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM) end # get_tmp_rec_(x::PAT.DiffCache{<:AbstractArray}, template) = PAT.get_tmp(x, template) # function get_tmp_rec_(x::Union{Tuple,NamedTuple}, template)