From e01b77a8b1ccdfbfaac9323c7b597b73ddd70a6b Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Wed, 28 Oct 2026 16:26:20 +0000 Subject: [PATCH 1/2] test point fit goes towards towards lower zetaM because logjac is lower. Compensated by higher logjac for zetaP, but zetaP is counted only once --- src/AbstractHybridProblem.jl | 21 +++--- src/HybridSolver_site.jl | 110 ++++++++++++++++++++++++++++++-- src/ModelApplicator.jl | 20 +++++- src/bijectors_utils.jl | 13 ++-- src/elbo_site.jl | 2 + src/elbo_site_grad.jl | 41 +++++++++++- test/test_elbo_site_doubleMM.jl | 96 ++++++++++++++++++++++++++-- 7 files changed, 275 insertions(+), 28 deletions(-) diff --git a/src/AbstractHybridProblem.jl b/src/AbstractHybridProblem.jl index 2c7cec5..28cae5b 100644 --- a/src/AbstractHybridProblem.jl +++ b/src/AbstractHybridProblem.jl @@ -210,8 +210,7 @@ function get_hybridproblem_transforms end Provide the number of covariates. Default implementation falls back to train_dataloader """ function get_hybridproblem_n_covar(prob::AbstractHybridProblem; scenario) - train_dataloader = get_hybridproblem_train_dataloader( - Random.default_rng(), prob; scenario) + train_dataloader = get_hybridproblem_train_dataloader(prob; scenario) (xM, xP, y_o, y_unc) = first(train_dataloader) n_covar = size(xM, 1) return (n_covar) @@ -277,16 +276,16 @@ function construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHyb #gdev = :use_gpu ∈ scenario ? gpu_device() : identity, ) (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; scenario) - i_sites_test = get_i_sites_test(prob; scenario) - n_site = size(xM, 2) - @assert size(xP,2) == n_site - @assert size(y_o,2) == n_site - @assert size(y_unc,2) == n_site - i_sites = setdiff(1:n_site, i_sites_test) - itrain_sites = 1:length(i_sites) # index into training data + i_indiv_test = get_i_sites_test(prob; scenario) + n_indiv = size(xM, 2) + @assert size(xP,2) == n_indiv + @assert size(y_o,2) == n_indiv + @assert size(y_unc,2) == n_indiv + i_indiv = setdiff(1:n_indiv, i_indiv_test) + i_indiv_train = 1:length(i_indiv) # index into training data train_loader = MLUtils.DataLoader(( - CA.getdata(xM)[:,i_sites], CA.getdata(xP)[:,i_sites], - y_o[:,i_sites], y_unc[:,i_sites], itrain_sites); + CA.getdata(xM)[:,i_indiv], CA.getdata(xP)[:,i_indiv], + y_o[:,i_indiv], y_unc[:,i_indiv], i_indiv_train); batchsize = n_batch, partial = false) return (train_loader) end diff --git a/src/HybridSolver_site.jl b/src/HybridSolver_site.jl index f06ba06..b3ccb0d 100644 --- a/src/HybridSolver_site.jl +++ b/src/HybridSolver_site.jl @@ -18,13 +18,13 @@ function HybridPointSolverSite(;alg = Optimisers.Adam(), n_MC_val=Val(3)) end function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolverSite{NMC_ELBO}, - loss_mode::AbstractElboLossMode; + loss_mode::AbstractElboLossMode, ϕ0; scenario=Val(()), rng=Random.default_rng(), epochs, epochs_callback = max(1, floor(Int, epochs/10)), -) where NMC_ELBO + show_train_loss::Val{TRAIN} = Val(false), +) where {NMC_ELBO, TRAIN} NMC = loss_mode isa NegLogDensity ? 1 : NMC_ELBO - loss_mode = NegLogDensity() train_loader = get_hybridproblem_train_dataloader(prob; scenario) # TODO replace key itrain_sites by i_indiv_train tmp = get_hybridproblem_test_data(prob; scenario) @@ -73,6 +73,11 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve h0 = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); h0_test = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site = n_indiv_test, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); + if TRAIN + indiv_args_train = NamedTuple{keys(indiv_args_test)}(train_loader.data) + n_indiv_train = size(indiv_args_train.xM,2) + h0_train = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site = n_indiv_train, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); + end rnormPM = prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=NMC, n_site) # py = get_hybridproblem_neg_logden_obs(prob; scenario) @@ -85,7 +90,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve transP, transM) sample_args_test = (;sample_args..., is_testmode = true) #nljoint_args_fix_test = (;nljoint_args_fix..., f! = ) - ϕ = (;ϕg = CA.getdata(ϕg), ϕqP, ϕqI) + ϕ_ref = (;ϕg = CA.getdata(ϕg), ϕqP, ϕqI) + ϕ = map(copy, ϕ0) # do not modify the initial guess # randnPM!(rng, rnormPM) #primal = neg_elbo_sites!(h0, loss_mode, ϕ, rnormPM, sample_args, indiv_args1, nljoint_args_fix) # res0, gradh0 = grad_neg_elbo_sites( @@ -94,6 +100,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve # ) # opt_state = Optimisers.setup(solver.alg, ϕ); # initialise this optimiser's momentum etc. + n_skips = 0 # epoch = 1 gradh = (;) stime = time() @@ -104,11 +111,15 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve indiv_args = NamedTuple{keys(indiv_args_test)}(data_batch_) #randnPM!(rng, rnormPM) # only required for NegElbo() grad_cv, gradh = grad_neg_elbo_sites( + #grad_cv, gradh = grad_neg_elbo_sites_enzyme( h0, gradh, loss_mode, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; solver.n_MC_val, #n_workers = 3, ) + #grad_cv.ϕqP .= zero(eltype(grad_cv.ϕqP)) # TODO remove fixing P gradient if any(map(x -> any(isnan.(x)), grad_cv)) + @show grad_cv.ϕqP, grad_cv.ϕqI + error("encountered non-finite gradient grad_v") n_skips += 1 #println("Skipped NaN : Batch $i") print(",$i") @@ -121,8 +132,99 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve # compute loss for test data loss_test = neg_elbo_sites!(h0_test, loss_mode, ϕ, rnormPM, sample_args_test, indiv_args_test, nljoint_args_fix) @show round(ttime, digits=1), epoch, loss_test + #@show exp.(ϕ.ϕqP[1:2]), grad_cv.ϕqP + if TRAIN + loss_train = neg_elbo_sites!(h0_train, loss_mode, ϕ, rnormPM, sample_args_test, indiv_args_train, nljoint_args_fix) + @show loss_train + end end end (;ϕ) end +function predict_point!(elbo_helpers, prob::AbstractHybridProblem, ϕ; scenario + ,keys_predict_indiv = (:ζsM, :θsM, :y)) + h0 = elbo_helpers + train_loader = get_hybridproblem_train_dataloader(prob; scenario) + keys_indiv = (:xM, :xP, :y_o, :y_unc, :i_indiv_train) + indiv_args1 = NamedTuple{keys_indiv}(first(train_loader)) + # TODO replace n_site by n_indiv + n_indiv_train, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario) + n_obs, n_site = size(indiv_args1.y_o) + # + par_templates = get_hybridproblem_par_templates(prob; scenario) + n_θP = length(par_templates.θP) + n_θM = length(par_templates.θM) + g, ϕg = get_hybridproblem_MLapplicator(prob; scenario) + # TODO: separate parameters from problem description - right now optimized + (; transP, transM) = get_hybridproblem_transforms(prob; scenario) + cor_ends = get_hybridproblem_cor_ends(prob; scenario) + approx1 = get_hybridproblem_HVIApproximation(prob; scenario) + # + # TODO make init_hybrid_ϕunc already return proper ϕqPc1 and ϕqIc1 + ϕ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), + logσ_ζP = ϕq.logσ_ζP, + ρsP = ϕq.ρsP, + ) + ϕqIc1 = CA.ComponentVector( + logσ_ζM_offsets = ϕq.logσ_ζM_offsets, + ρsM = ϕq.ρsM, + ) + ϕqP = ϕqP2 = CA.getdata(ϕqPc1) + ϕqI = CA.getdata(ϕqIc1) + intϕqP = get_concrete(ComponentArrayInterpreter(ϕqPc1)) + intϕqI = get_concrete(ComponentArrayInterpreter(ϕqIc1)) + # + diffchunk = ForwardDiff.Chunk(8) + pbm = create_nsite_applicator(get_hybridproblem_PBmodel(prob; scenario), 1; diffchunk) + nljoint_args_fix =(; + f! = pbm, + fLogLik = get_hybridproblem_neg_logden_obs(prob; scenario) + ) + n_M = get_numberof_MLinputs(approx, par_templates.θM) + n_cov = get_hybridproblem_n_covar(prob; scenario) + pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) + n_covP = length(pbm_covars) + rnormPM = prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=1, n_site) + # + py = get_hybridproblem_neg_logden_obs(prob; scenario) + # priors = get_hybridproblem_priors(prob; scenario) + # priorsP = Tuple(priors[k] for k in keys(par_templates.θP)) + # priorsM = Tuple(priors[k] for k in keys(par_templates.θM)) + pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_templates.θP)) for k in pbm_covars]) + sample_args = (; + approx, g, is_testmode = false, pbm_covar_indices, intϕqP, intϕqI, cor_ends, + transP, transM) + sample_args_test = (;sample_args..., is_testmode = true) + #nljoint_args_fix_test = (;nljoint_args_fix..., f! = ) + # convert from DiffCache to normal vector + h0p = if h0.helpers_sites[1][1] isa PAT.DiffCache + h0p = (; h0..., helpers_sites = map_leaves_nt(x -> PAT.get_tmp(x, par_templates.θM), h0.helpers_sites)) + else + h0 + end + @assert n_batch == size(indiv_args1.xM,2) + keys_predict_pop = (:ζsP, :θsP) + h1 = h0p.helpers_sites[1] + # do not know the number of training sites sum of sites in batch, need to + # reduce by vcat + hi_agg_batch = NamedTuple{keys_predict_indiv}(map(keys_predict_indiv) do key + repeat(h1[key][:,1], 1, n_batch) # extend columns to n_site_train + end) + red = make_tuple_reducer(hcat) + hi_agg = mapfoldl(red, enumerate(train_loader), init=map(x -> x[:,1:0],hi_agg_batch)) do (i_batch, data_batch) + indiv_args = NamedTuple{keys_indiv}(data_batch) + neg_elbo_sites!(h0p, NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + for key in keys_predict_indiv + for ii in 1:n_batch + hi_agg_batch[key][:,ii] .= h0p.helpers_sites[ii][key] + end + end + hi_agg_batch + end + h0p[keys_predict_pop], hi_agg +end + diff --git a/src/ModelApplicator.jl b/src/ModelApplicator.jl index 6a9a95e..3f3412f 100644 --- a/src/ModelApplicator.jl +++ b/src/ModelApplicator.jl @@ -236,6 +236,9 @@ function apply_model!(y, app::NormalScalingModelApplicator, x, ϕ; kwargs...) # @show typeof(app.μ) # @show typeof(ϕ) @assert eltype(app.μ) == eltype(ϕ) + #y[app.range_scaled,:] + #norminvcdf.(app.μ, app.σ, y[app.range_scaled,:]) + #norminvcdf.(app.μ, app.σ, tmp[app.range_scaled,:]) if !isnothing(app.range_scaled) for x in app.range_scaled; @assert(x ∈ axes(y,1)); end # .= fails with Enzyme @@ -284,7 +287,7 @@ end function apply_model(app::RangeScalingModelApplicator, x, ϕ; kwargs...) res0 = apply_model(app.app, x, ϕ; kwargs...) - if !isempty(app.range_scaled) + if !isnothing(app.range_scaled) && !isempty(app.range_scaled) res_scaled = index_firstdim(res0,app.range_scaled) .* app.width .+ app.offset combine_range(res0, res_scaled, app.range_scaled) else @@ -292,6 +295,21 @@ function apply_model(app::RangeScalingModelApplicator, x, ϕ; kwargs...) end end +function apply_model!(y, app::RangeScalingModelApplicator, x, ϕ; kwargs...) + apply_model!(y, app.app, x, ϕ; kwargs...) + @assert eltype(app.width) == eltype(ϕ) + if !isnothing(app.range_scaled) + for x in app.range_scaled; @assert(x ∈ axes(y,1)); end + @inbounds for j in axes(y,2), (i,ri) in enumerate(app.range_scaled) + y[ri,j] = y[ri,j] * app.width[i] + app.offset[i] + end + else + y .= y .* app.width .+ app.offset + end + y +end + + function combine_range(res0, res_scaled, range_scaled) range_before = 1:(range_scaled[1]-1) range_after = (range_scaled[end]+1):size(res0,1) diff --git a/src/bijectors_utils.jl b/src/bijectors_utils.jl index 7056bab..c2bcaa6 100644 --- a/src/bijectors_utils.jl +++ b/src/bijectors_utils.jl @@ -56,7 +56,9 @@ Bijectors.logabsdetjac(b::Exp, x) = sum(x) function Bijectors.with_logabsdet_jacobian(b::Exp, x) return exp.(x), sum(x) end -function Bijectors.with_logabsdet_jacobian!(b::Exp, x::AbstractArray{T}, y, logjac=zero(T)) where T +function Bijectors.with_logabsdet_jacobian!( + b::Exp, x::AbstractArray{T}, y, logjac=zero(T) +) where T # order of arguemts (y not first) given by Bijectors y .= exp.(x) # for i in eachindex(y) # did not decrease allocations with Enzyme @@ -228,8 +230,9 @@ function Bijectors.with_logabsdet_jacobian!(f::typeof(identity), x::AbstractArra end #------------------------- mutating variants -function with_logabsdet_jacobian_stacked!(y, - b::Stacked{<:NTuple{N,Any},<:NTuple{N,Any}}, x::AbstractVector{T}, logjac=zero(T)) where {N,T} +function with_logabsdet_jacobian_stacked!( + y, b::Stacked{<:NTuple{N,Any},<:NTuple{N,Any}}, x::AbstractVector{T}, logjac=zero(T) +) where {N,T} for i in 1:N local rin = b.ranges_in[i] local rout = b.ranges_out[i] @@ -240,9 +243,11 @@ function with_logabsdet_jacobian_stacked!(y, (y, logjac) end -function with_logabsdet_jacobian_stacked!(ys::AbstractMatrix{T}, bs::Stacked, xs::AbstractMatrix{T}, logjac = zero(T)) where T +function with_logabsdet_jacobian_stacked!( + ys::AbstractMatrix{T}, bs::Stacked, xs::AbstractMatrix{T}, logjac = zero(T)) where T @assert size(ys) == size(xs) @inbounds for j in 1:size(xs,2) + # logjac is already the sum of previous (provided to method) and the transformation (_, logjac) = with_logabsdet_jacobian_stacked!(view(ys,:,j), bs, view(xs,:,j), logjac) end (ys, logjac) diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 8a21430..12d18d5 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -69,6 +69,7 @@ function neg_elbo_sites!( # loglik = sum(x -> x.loglik, res_site) # costTrans = sum(x -> x.costTrans, res_site) #elbo = sum(first, res_site) - sum(h.logσ_ζP) + error("inpsect elbo components") elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) end @@ -107,6 +108,7 @@ function compute_nelboi_z!(hi, loss_mode::AbstractElboLossMode, approx::Abstract 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, hi)[1] + @show exp_nL, exp_ladJacTM, sum(hi.logσ_ζM) nelboi_z = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM) end diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 7568571..6e980d2 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -1,3 +1,38 @@ +function _ftmp_point(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) + neg_elbo_sites!(h0, NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] +end +function grad_neg_elbo_sites_enzyme( + elbo_helpers::NamedTuple, # tuple of preallocated arrays + grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures + loss_mode::NegLogDensity, + ϕ::NamedTuple, + rnormPM::NamedTuple, + sample_args::NamedTuple, + indiv_args::NamedTuple, + nljoint_args_fix::NamedTuple; + n_MC_val::Val = Val(3), + n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... + executor::Transducers.Executor = Transducers.SequentialEx(), +) + h = elbo_helpers + rnormPM_point = (;rnormPM.P, M = map(x -> x[:,1:1], rnormPM.M)) + sample_args_rev = (; sample_args..., is_testmode=false) + (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args_rev + primal = _ftmp_point(h, ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) + dϕ = map(zero,ϕ) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse) , + _ftmp_point, + Enzyme.Active, + Enzyme.Duplicated(h, Enzyme.make_zero(h)), + Enzyme.Duplicated(ϕ, dϕ), + Enzyme.DuplicatedNoNeed(rnormPM_point, Enzyme.make_zero(rnormPM_point)), + Enzyme.Const(sample_args_rev), + Enzyme.Const(indiv_args), + Enzyme.Const(nljoint_args_fix), + ) + dϕ, h +end function grad_neg_elbo_sites( elbo_helpers::NamedTuple, # tuple of preallocated arrays @@ -30,7 +65,7 @@ function grad_neg_elbo_sites( NMC = NMC_ELBO end h = elbo_helpers # preallocated μζP, dμζP, ζsP, ϕms, xMP, dxMP - (;ϕg, ϕqP, ϕqI) = ϕ + (;ϕg, ϕqP, ϕqI) = map(CA.getdata, ϕ) TF = eltype(ϕqP) (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args @assert keys(indiv_args)[1] == :xM @@ -61,6 +96,7 @@ function grad_neg_elbo_sites( # set_ζsP!(h.ζsP, h.logσ_ζP, loss_mode, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers) # n_P * n_MC g_apply!(h_ϕm, ϕg, xM, h.ζsP, pbm_covar_indices, g, h.xMP, is_testmode) + @assert all(isfinite.(h_ϕm)) ladJacTP = transformζ!(h.θsP, transP, h.ζsP) # return value captures ladJacT # # parallel ForwardDiffGradient through forwarddiff_grad_nelboi_z! @@ -121,6 +157,9 @@ function grad_neg_elbo_sites( h.ζsP, h.logσ_ζP, approx, rnormPM.P, ϕqPc, cor_ends.P, h.sample_buffers, ) # need same names as in ϕ for Optimizer + if !all(isfinite.(gradh.dϕg)) || !all(isfinite.(gradh.dϕqPc)) || !all(isfinite.(gradh.dϕqIc)) + error("encountered noninfite gradient") + end NamedTuple{keys(ϕ)}((gradh.dϕg, gradh.dϕqPc, gradh.dϕqIc, )), gradh end diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 5b9da76..3e2feee 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -23,16 +23,19 @@ import ForwardDiff import Optimisers using CommonSolve +import MLUtils +import PreallocationTools as PAT + -rng = StableRNG(1234) prob = DoubleMM.DoubleMMCase() scenario = Val((:use_Lux,:usePopPBM,:covarK2,:uniscaling,)) #scenario = Val((:use_Lux,:usePopPBM,:uniscaling,)) # omit covars -> vector site ϕm - +#scenario = Val((:use_Lux,:usePopPBM,:uniscaling,:use_rangescaling)) # omit covars, rangescaling diffchunk = ForwardDiff.Chunk(8) +rng = StableRNG(1234) (; 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) @@ -57,8 +60,10 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te # TODO update init_hybrid_ϕunc ϕq, approx = init_hybrid_ϕunc(approx1, cor_ends; transM, par_templates.θM, n_site) + μ0_bias = 0.9 + μ0_bias = 1.0 # true μP ϕqPc1 = CA.ComponentVector( - μζP = log.(CA.getdata(par_templates.θP)), + μζP = log.(CA.getdata(par_templates.θP .* eltype(par_templates.θP)(μ0_bias))), #logσ_ζP = ones(n_θP) .* log(0.01), logσ_ζP = ϕq.logσ_ζP, ρsP = ϕq.ρsP, @@ -92,6 +97,7 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te randn!(h0.θsP) # for testing should initialized to finite values randn!(h0_point.θsP) # for testing should initialized to finite values rnormPM = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC, n_site) + rnormPM_point = CP.prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=1, n_site) #------------------------------------------------------------- @testset "grad_neg_elbo_sites" begin @@ -158,12 +164,88 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te alloc_grad_neg_elbo_sites(h0_point, gradh0_point, CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) end +@testset "grad_neg_elbo_sites_enzyme" begin + local ϕqIc, ϕqPc, rng1, indiv_args, ϕ, sample_args, primal + ϕqIc = intϕqI(ϕqI) + ϕqPc = intϕqP(ϕqP) + # + 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 = true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, + transP, transM) + primal = CP.neg_elbo_sites!(h0_point, CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) + dϕ, _ = CP.grad_neg_elbo_sites_enzyme( + h0_point, (;), CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix, + ) + dϕ.ϕqP +end + @testset "HybridPointSolverSite" begin - solver_point = HybridPointSolverSite() - res = CommonSolve.solve(prob, solver_point, CP.NegLogDensity(); scenario, epochs=2) - + # in order to test inversion, set the dataloader to the synthetic data + rng1 = rng = StableRNG(111) + local xM, xP, y_o, y_unc, θP_true, θMs_true, y_true + (; xM, xP, y_o, y_unc, θP_true, θMs_true, y_true) = HVI.gen_hybridproblem_synthetic(rng1, prob; scenario) + # constrain to training dataset + i_indiv_test = get_i_sites_test(prob; scenario) + i_indiv = setdiff(1:n_indiv, i_indiv_test) + (; xM, xP, y_o, y_unc, θMs_true, y_true) = map( + (; xM, xP, y_o, y_unc, θMs_true, y_true)) do x + selectdim(x,ndims(x),i_indiv) + end + # local n_indiv, i_indiv_test, i_indiv, i_indiv_train, n_batch, train_dataloader, test_data + # n_indiv = size(xM, 2) + # i_indiv_test = HVI.get_i_sites_test(prob; scenario) + # i_indiv = setdiff(1:n_indiv, i_indiv_test) + # i_indiv_train = 1:length(i_indiv) # index into training data + # n_batch = length(h0_point.helpers_sites) + # train_dataloader = MLUtils.DataLoader(( + # CA.getdata(xM)[:,i_indiv], CA.getdata(xP)[:,i_indiv], + # y_o[:,i_indiv], y_unc[:,i_indiv], i_indiv_train); + # batchsize = n_batch, partial = false) + # test_data = (; xM=xM[:,i_indiv_test], xP=xP[:,i_indiv_test], y_o=y_o[:,i_indiv_test], y_unc=y_unc[:,i_indiv_test], itrain_sites=zeros(length(i_indiv_test))) + # prob_syn = HybridProblem(prob; scenario, train_dataloader, test_data) + prob_syn = prob + y_o2 = get_hybridproblem_train_dataloader(prob; scenario).data[3] + y_o .- y_o2 + solver_point = HybridPointSolverSite(alg=Optimisers.Adam(0.0001)) + ϕ0 = (;ϕg = copy(CA.getdata(ϕg)), ϕqP=copy(ϕqP), ϕqI=copy(ϕqI)) + exp.(ϕ0.ϕqP[1:2]) # updated during solve + #ϕ0 = res.ϕ + res = CommonSolve.solve(prob_syn, solver_point, CP.NegLogDensity(), ϕ0; scenario, + epochs=20000, + show_train_loss = Val(true) + ) + exp.(res.ϕ.ϕqP[1:2]) + #@test exp.(intϕqP(res.ϕ.ϕqP).μζP) ≈ θP_true atol = 0.1 + # run forward + ϕ_tmp = (;res.ϕ.ϕg, ϕ0.ϕqP, ϕ0.ϕqI) + hp, hi_agg = HVI.predict_point!(h0_point, prob_syn, res.ϕ; scenario) + hp0, hi_agg0 = HVI.predict_point!(h0_point, prob_syn, ϕ0; scenario) + hpt, hi_aggt = HVI.predict_point!(h0_point, prob_syn, ϕ_tmp; scenario) + @test hp.θsP[:,1] ≈ θP_true atol=0.01 + () -> begin + # @usingany UnicodePlots + i_obs = 3 + #scatterplot(hi_agg.θsM[i_obs,:], θMs_true[i_obs,:]) + scatterplot(y_true[i_obs,:], hi_agg.y[i_obs,:]) + scatterplot(hi_agg0.y[i_obs,:], y_true[i_obs,:]) + scatterplot(hi_aggt.y[i_obs,:], y_true[i_obs,:]) + scatterplot(hi_aggt.y[i_obs,:], hi_agg.y[i_obs,:]) + sum(abs2, hi_agg0.y .- y_true) + sum(abs2, hi_agg.y .- y_true) + sum(abs2, hi_aggt.y .- y_true) # with true ζP much worse? + scatterplot(θMs_true[1,:], hi_agg.θsM[1,:], ) + scatterplot(θMs_true[2,:], hi_agg.θsM[2,:]) + end + @test hi_agg.θsM ≈ θMs_true atol=0.01 + h0p.helpers_sites[1].θsM #@test res.ϕ != ϕ - + x = h0p.helpers_sites[1] end From af4fe0ff3d841f9ed554907a086e54113e74d7f7 Mon Sep 17 00:00:00 2001 From: Thomas Wutzler Date: Thu, 29 Oct 2026 15:59:49 +0000 Subject: [PATCH 2/2] devide exp_ladJacTM by n_indiv to prevent overly large influence of transformation with many sites. --- src/HybridSolver_site.jl | 24 +++++--- src/elbo_site.jl | 9 ++- src/elbo_site_grad.jl | 79 ++++++++++++++------------- test/test_elbo_site.jl | 49 +++++++++-------- test/test_elbo_site_doubleMM.jl | 97 +++++++++++++++------------------ 5 files changed, 127 insertions(+), 131 deletions(-) diff --git a/src/HybridSolver_site.jl b/src/HybridSolver_site.jl index b3ccb0d..b2549f2 100644 --- a/src/HybridSolver_site.jl +++ b/src/HybridSolver_site.jl @@ -37,8 +37,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve par_templates = get_hybridproblem_par_templates(prob; scenario) n_θP = length(par_templates.θP) n_θM = length(par_templates.θM) - g, ϕg = get_hybridproblem_MLapplicator(prob; scenario) # TODO: separate parameters from problem description - right now optimized + g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario) (; transP, transM) = get_hybridproblem_transforms(prob; scenario) cor_ends = get_hybridproblem_cor_ends(prob; scenario) approx1 = get_hybridproblem_HVIApproximation(prob; scenario) @@ -55,10 +55,16 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve logσ_ζM_offsets = ϕq.logσ_ζM_offsets, ρsM = ϕq.ρsM, ) - ϕqP = ϕqP2 = CA.getdata(ϕqPc1) - ϕqI = CA.getdata(ϕqIc1) intϕqP = get_concrete(ComponentArrayInterpreter(ϕqPc1)) intϕqI = get_concrete(ComponentArrayInterpreter(ϕqIc1)) + if ϕ0 == NamedTuple() + ϕqP = CA.getdata(ϕqPc1) + ϕqI = CA.getdata(ϕqIc1) + ϕ0 = (;ϕg = CA.getdata(ϕg0), ϕqP, ϕqI) # copied - not modified + else + @assert length(ϕ0.ϕqP) == length(ϕqPc1) + @assert length(ϕ0.ϕqI) == length(ϕqIc1) + end # diffchunk = ForwardDiff.Chunk(8) pbm = create_nsite_applicator(get_hybridproblem_PBmodel(prob; scenario), NMC; diffchunk) @@ -71,14 +77,14 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario) n_covP = length(pbm_covars) - h0 = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); - h0_test = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site = n_indiv_test, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); + h0 = prepare_elbo_helpers(approx, ϕ0.ϕg, ϕ0.ϕqP; n_θP, n_θM, n_site, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); + h0_test = prepare_elbo_helpers(approx, ϕ0.ϕg, ϕ0.ϕqP; n_θP, n_θM, n_site = n_indiv_test, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); if TRAIN indiv_args_train = NamedTuple{keys(indiv_args_test)}(train_loader.data) n_indiv_train = size(indiv_args_train.xM,2) - h0_train = prepare_elbo_helpers(approx, ϕg, ϕqP; n_θP, n_θM, n_site = n_indiv_train, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); + h0_train = prepare_elbo_helpers(approx, ϕ0.ϕg, ϕ0.ϕqP; n_θP, n_θM, n_site = n_indiv_train, n_MC=NMC, n_cov, n_covP, n_M, cor_ends, n_obs); end - rnormPM = prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=NMC, n_site) + rnormPM = prepare_rnorm(ϕ0.ϕqP; n_θP, n_θM, n_MC=NMC, n_site) # py = get_hybridproblem_neg_logden_obs(prob; scenario) # priors = get_hybridproblem_priors(prob; scenario) @@ -90,7 +96,6 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve transP, transM) sample_args_test = (;sample_args..., is_testmode = true) #nljoint_args_fix_test = (;nljoint_args_fix..., f! = ) - ϕ_ref = (;ϕg = CA.getdata(ϕg), ϕqP, ϕqI) ϕ = map(copy, ϕ0) # do not modify the initial guess # randnPM!(rng, rnormPM) #primal = neg_elbo_sites!(h0, loss_mode, ϕ, rnormPM, sample_args, indiv_args1, nljoint_args_fix) @@ -139,7 +144,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve end end end - (;ϕ) + (;ϕ, h0, h0_test) end function predict_point!(elbo_helpers, prob::AbstractHybridProblem, ϕ; scenario @@ -207,6 +212,7 @@ function predict_point!(elbo_helpers, prob::AbstractHybridProblem, ϕ; scenario h0 end @assert n_batch == size(indiv_args1.xM,2) + # keys_predict_pop = (:ζsP, :θsP) h1 = h0p.helpers_sites[1] # do not know the number of training sites sum of sites in batch, need to diff --git a/src/elbo_site.jl b/src/elbo_site.jl index 12d18d5..fd0b32c 100644 --- a/src/elbo_site.jl +++ b/src/elbo_site.jl @@ -53,12 +53,13 @@ function neg_elbo_sites!( template = ϕqI # only important for gradient θsP = h.θsP keys_nljoint_args_indi_val = Val(keys(nljoint_args_inds)) + n_indiv = size(h.ϕms,2) # closure with approx, transM, and kwargs function compute_nelboi_z_cl!(hi, rnormM, nljoint_args_indi_tup, ϕm) #nljoint_args_indi = NamedTuple{keys_nljoint_args_indi}(nljoint_args_indi_tup) nljoint_args_indi = name_tuple(nljoint_args_indi_tup, keys_nljoint_args_indi_val) compute_nelboi_z!(hi, loss_mode, approx, rnormM, nljoint_args_indi, ϕm, - ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) + ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix, n_indiv) end #res_site = map(compute_nelboi_z!, loss_mode, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it) #MAYBE: distributed mapreduce: @@ -69,7 +70,6 @@ function neg_elbo_sites!( # loglik = sum(x -> x.loglik, res_site) # costTrans = sum(x -> x.costTrans, res_site) #elbo = sum(first, res_site) - sum(h.logσ_ζP) - error("inpsect elbo components") elbo = elbo_z - exp_ladJacTP - sum(h.logσ_ζP) end @@ -94,7 +94,7 @@ name_tuple(tup, ::Val{KEYS}) where KEYS = NamedTuple{KEYS}(tup) function compute_nelboi_z!(hi, loss_mode::AbstractElboLossMode, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, ϕm::AbstractArray, ϕqIc::AbstractArray, θsP, - cor_endsM, transM, nljoint_args_fix, ) + cor_endsM, transM, nljoint_args_fix, n_indiv::Int) # DM is 1 (vector) if no population covariates - constant across MC # 2 (Matrix) if g is driven by population covariates - differs across MC # on update -> sync corresponding function within grad_neg_elbo_sites @@ -108,8 +108,7 @@ function compute_nelboi_z!(hi, loss_mode::AbstractElboLossMode, approx::Abstract 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, hi)[1] - @show exp_nL, exp_ladJacTM, sum(hi.logσ_ζM) - nelboi_z = exp_nL - exp_ladJacTM - sum(hi.logσ_ζM) + nelboi_z = exp_nL - (exp_ladJacTM ./ n_indiv) - sum(hi.logσ_ζM) end @inline function set_ζsM!(ζsM, logσ_ζM, ::NegElbo, approx, rnormM, ϕqIc, ϕm, cor_endsM, sample_buffers) diff --git a/src/elbo_site_grad.jl b/src/elbo_site_grad.jl index 6e980d2..bdc48b5 100644 --- a/src/elbo_site_grad.jl +++ b/src/elbo_site_grad.jl @@ -1,38 +1,38 @@ -function _ftmp_point(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) - neg_elbo_sites!(h0, NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] -end -function grad_neg_elbo_sites_enzyme( - elbo_helpers::NamedTuple, # tuple of preallocated arrays - grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures - loss_mode::NegLogDensity, - ϕ::NamedTuple, - rnormPM::NamedTuple, - sample_args::NamedTuple, - indiv_args::NamedTuple, - nljoint_args_fix::NamedTuple; - n_MC_val::Val = Val(3), - n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... - executor::Transducers.Executor = Transducers.SequentialEx(), -) - h = elbo_helpers - rnormPM_point = (;rnormPM.P, M = map(x -> x[:,1:1], rnormPM.M)) - sample_args_rev = (; sample_args..., is_testmode=false) - (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args_rev - primal = _ftmp_point(h, ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) - dϕ = map(zero,ϕ) - Enzyme.autodiff( - Enzyme.set_runtime_activity(Enzyme.Reverse) , - _ftmp_point, - Enzyme.Active, - Enzyme.Duplicated(h, Enzyme.make_zero(h)), - Enzyme.Duplicated(ϕ, dϕ), - Enzyme.DuplicatedNoNeed(rnormPM_point, Enzyme.make_zero(rnormPM_point)), - Enzyme.Const(sample_args_rev), - Enzyme.Const(indiv_args), - Enzyme.Const(nljoint_args_fix), - ) - dϕ, h -end +# function _ftmp_point(h0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix) +# neg_elbo_sites!(h0, NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix)[1] +# end +# function grad_neg_elbo_sites_enzyme( +# elbo_helpers::NamedTuple, # tuple of preallocated arrays +# grad_elbo_helpers::NamedTuple, # tuple of preallocated arrays pullback closures +# loss_mode::NegLogDensity, +# ϕ::NamedTuple, +# rnormPM::NamedTuple, +# sample_args::NamedTuple, +# indiv_args::NamedTuple, +# nljoint_args_fix::NamedTuple; +# n_MC_val::Val = Val(3), +# n_workers = 1, # TODO compute from executor and Distributed.nworkers and nthreads,... +# executor::Transducers.Executor = Transducers.SequentialEx(), +# ) +# h = elbo_helpers +# rnormPM_point = (;rnormPM.P, M = map(x -> x[:,1:1], rnormPM.M)) +# sample_args_rev = (; sample_args..., is_testmode=false) +# (;approx, g, is_testmode, pbm_covar_indices, intϕqP, intϕqI, cor_ends, transP, transM) = sample_args_rev +# primal = _ftmp_point(h, ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) +# dϕ = map(zero,ϕ) +# Enzyme.autodiff( +# Enzyme.set_runtime_activity(Enzyme.Reverse) , +# _ftmp_point, +# Enzyme.Active, +# Enzyme.Duplicated(h, Enzyme.make_zero(h)), +# Enzyme.Duplicated(ϕ, dϕ), +# Enzyme.DuplicatedNoNeed(rnormPM_point, Enzyme.make_zero(rnormPM_point)), +# Enzyme.Const(sample_args_rev), +# Enzyme.Const(indiv_args), +# Enzyme.Const(nljoint_args_fix), +# ) +# dϕ, h +# end function grad_neg_elbo_sites( elbo_helpers::NamedTuple, # tuple of preallocated arrays @@ -217,12 +217,13 @@ function forwarddiff_grad_nelboi_z!(hi, rnormM, # group arrays of same dimensionality (here grad_ϕm and grad_ϕqI) into a Vcat ax_ϕqI = CA.getaxes(cl.ϕqIc)[1] ax_ϕP = get_axis_θP(cl.nljoint_args_fix.f!) + n_indiv = size(cl.dϕmvecs,2) with_channel_element(cl.hw_channel) do hwi_tup local hwi, hwi_dc, nelboi_vv, θsP_ϕm_ϕqI_vv, grad_θsP, grad_ϕm, grad_ϕqI hwi, hwi_dc = hwi_tup #grad_conf = hwi.grad_conf nelboi_vv = make_nelboiz_lazy_cls_vv(hi, cl.approx, rnormM, nljoint_args_indi, - cl.cor_endsM, cl.transM, cl.nljoint_args_fix, cl.loss_mode, ax_ϕqI) + cl.cor_endsM, cl.transM, cl.nljoint_args_fix, n_indiv, cl.loss_mode, ax_ϕqI) # grads are vector-shaped, but copyto! dϕmvecs and SAVector use those anyhow grad_θsP, grad_ϕm, grad_ϕqI = if isnothing(omit_gradient) # θsP_ϕm_ϕqI is preallocated per worker (hwi.θsP_ϕm_ϕqI). @@ -253,7 +254,7 @@ end # TODO remove unused arguments function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation, rnormM, nljoint_args_indi, - cor_endsM, transM, nljoint_args_fix, + cor_endsM, transM, nljoint_args_fix, n_indiv, loss_mode::AbstractElboLossMode, ax_ϕqI::CA.Axis ) @@ -267,7 +268,7 @@ function make_nelboiz_lazy_cls_vv(hi, approx::AbstractHVIApproximation, compute_nelboi_z!( hi, loss_mode, approx, rnormM, nljoint_args_indi, ϕm_dual, ϕqIc_dual, θsP_dual, - cor_endsM, transM, nljoint_args_fix, + cor_endsM, transM, nljoint_args_fix, n_indiv, )[1] end end @@ -479,7 +480,7 @@ function prepare_gradelbo_helpers( # local nelboiz_vv, θsP_ϕm_ϕqI_vv1, diffchunks_elbo_vv nelboiz_vv = make_nelboiz_lazy_cls_vv(hi1, approx, rnormMi, nljoint_args_indi, - cor_ends.M, transM, nljoint_args_fix, loss_mode, CA.getaxes(ϕqIc)[1], ) + cor_ends.M, transM, nljoint_args_fix, n_site, loss_mode, CA.getaxes(ϕqIc)[1], ) θsP_ϕm_ϕqI_vv1 = θsP_ϕm_ϕqI_vv(θsP,ϕm,ϕqIc) # n_M >= n_θM but not static diffchunks_elbo_vv = ForwardDiff.Chunk(min(chunk, diff --git a/test/test_elbo_site.jl b/test/test_elbo_site.jl index 5872d64..17b54e4 100644 --- a/test/test_elbo_site.jl +++ b/test/test_elbo_site.jl @@ -419,17 +419,17 @@ end ϕms1 = h0.ϕms[:,1] θsP1 = h0.θsP CP.compute_nelboi_z!(hi1, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, - cor_ends.M, transM, nljoint_args_fix) + cor_ends.M, transM, nljoint_args_fix, n_site) #@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, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) + function tmpf(hi, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix, n_site) @test (@allocated CP.compute_nelboi_z!(hi, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1, - ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 + ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix, n_site)) == 0 end - tmpf(hi1, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) + tmpf(hi1, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix, n_site) CP.compute_nelboi_z!(hi1, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, - cor_ends.M, transM, nljoint_args_fix) - tmpf(hi1, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) + cor_ends.M, transM, nljoint_args_fix, n_site) + tmpf(hi1, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix, n_site) # # test with views as input to compute_nelboi_z! and differnt ϕm per MC h21 = h2.helpers_sites[1] @@ -437,25 +437,25 @@ end ϕm = randn!(similar(h2.ϕms_mcs[:,:,1])) inputs = CA.ComponentVector(ϕqIc=ϕqIc, ϕm = ϕm, θsP= θsP1) CP.compute_nelboi_z!(h21, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, inputs.ϕm, inputs.ϕqIc, - inputs.θsP, cor_ends.M, transM, nljoint_args_fix) + inputs.θsP, cor_ends.M, transM, nljoint_args_fix, n_site) # with ϕms1_ = view(inputs, Val(:ϕm)) ϕqIc_ = view(inputs, Val(:ϕqIc)) θsP1_ = view(inputs, Val(:θsP)) - function loop_compute_nelboi_z(n, h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + function loop_compute_nelboi_z(n, h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site) for _ in 1:n - @noinline CP.compute_nelboi_z!(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + @noinline CP.compute_nelboi_z!(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site) end nothing end - function alloc_compute_nelboi_z(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) + function alloc_compute_nelboi_z(h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site) # avoid global variables -> pass them through function - loop_compute_nelboi_z(1, h21, loss_mode, 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, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix) - @test (@allocated loop_compute_nelboi_z(100,h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix)) == 0 + loop_compute_nelboi_z(1, h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site) # thorough warmup + #@profview_allocs loop_compute_nelboi_z(10_000,h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site) + @test (@allocated loop_compute_nelboi_z(100,h21, loss_mode, approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_endsM, transM, nljoint_args_fix, n_site)) == 0 end - alloc_compute_nelboi_z(h21, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) - alloc_compute_nelboi_z(h21, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix) + alloc_compute_nelboi_z(h21, CP.NegElbo(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix, n_site) + alloc_compute_nelboi_z(h21, CP.NegLogDensity(), approx, rnormM1, nljoint_args_ind1, ϕms1_, ϕqIc_, θsP1_, cor_ends.M, transM, nljoint_args_fix, n_site) # gradh2 = @inferred CP.prepare_gradelbo_helpers(inputs.ϕqIc, inputs.ϕm, inputs.θsP, ϕg, ϕqPc, approx; pbm_covar_indices=pbm_covar_indices2, n_workers=1, @@ -497,14 +497,14 @@ end approx2 = MeanHVIApproximation() hiM1 = h0M.helpers_sites[1] @test isfinite(CP.compute_nelboi_z!(hiM1, CP.NegElbo(), approx2, rnormM1, nljoint_args_ind1, - ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix)) + ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix, n_site)) function alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, - ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix) + ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix, n_site) @test (@allocated CP.compute_nelboi_z!(hiM1, CP.NegElbo(), approx2, rnormM1, nljoint_args_ind1, - ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix)) == 0 + ϕms1, ϕqIc, θsP1, cor_endsM, transM, nljoint_args_fix, n_site)) == 0 end alloc_compute_nelboi_zM!(hiM1, approx2, rnormM1, nljoint_args_ind1, - ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix) + ϕms1, ϕqIc, θsP1, cor_ends.M, transM, nljoint_args_fix, n_site) end @testset "pullback_g_apply!" begin @@ -994,7 +994,7 @@ end primal2S = CP.neg_elbo_sites!(h2S, CP.NegElbo(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S) grad0, gradh2S = CP.grad_neg_elbo_sites( h2S, (;), CP.NegElbo(), ϕ2S, rnormPM, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, - #executor = distributedEx, + executor = distributedEx, ) if isfile("intermediate/test_enzymeS_dphi2.jld2") primal_enz2S, dϕg2S_enz, dϕqP2S_enz, dϕqI2S_enz = JLD2.load( @@ -1008,14 +1008,15 @@ end #hcat(dϕqI2S_enz, grad0.dϕqI) end # + ϕ2S_point = map(copy, ϕ2S) rng1 = StableRNG(1234) CP.randnPM!(rng1, rnormPM_point) - randn!(rng1, ϕ2S.ϕg) + randn!(rng1, ϕ2S_point.ϕg) randn!(rng1, indiv_args.xM) - primal2S_point = CP.neg_elbo_sites!(h2S_point, CP.NegLogDensity(), ϕ2S, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S) + primal2S_point = CP.neg_elbo_sites!(h2S_point, CP.NegLogDensity(), ϕ2S_point, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S) grad_point, gradh2S_point = CP.grad_neg_elbo_sites( - h2S_point, (;), CP.NegLogDensity(), ϕ2S, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, - #executor = distributedEx, + h2S_point, (;), CP.NegLogDensity(), ϕ2S_point, rnormPM_point, sample_args2S, indiv_args, nljoint_args_fix2S; n_MC_val, + executor = distributedEx, ) if isfile("intermediate/test_enzymeS_dphi2_point.jld2") primal_enz2S_point, dϕg2S_point_enz, dϕqP2S_point_enz, dϕqI2S_point_enz = JLD2.load( diff --git a/test/test_elbo_site_doubleMM.jl b/test/test_elbo_site_doubleMM.jl index 3e2feee..acd7969 100644 --- a/test/test_elbo_site_doubleMM.jl +++ b/test/test_elbo_site_doubleMM.jl @@ -29,9 +29,9 @@ import PreallocationTools as PAT prob = DoubleMM.DoubleMMCase() -scenario = Val((:use_Lux,:usePopPBM,:covarK2,:uniscaling,)) +#scenario = Val((:use_Lux,:usePopPBM,:covarK2,:uniscaling,)) #scenario = Val((:use_Lux,:usePopPBM,:uniscaling,)) # omit covars -> vector site ϕm -#scenario = Val((:use_Lux,:usePopPBM,:uniscaling,:use_rangescaling)) # omit covars, rangescaling +scenario = Val((:use_Lux,:usePopPBM,:uniscaling,:use_rangescaling)) # omit covars, rangescaling diffchunk = ForwardDiff.Chunk(8) @@ -164,26 +164,26 @@ pbm_covar_indices = SA.SVector{length(pbm_covars)}([findfirst(==(k), keys(par_te alloc_grad_neg_elbo_sites(h0_point, gradh0_point, CP.NegLogDensity(), ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) end -@testset "grad_neg_elbo_sites_enzyme" begin - local ϕqIc, ϕqPc, rng1, indiv_args, ϕ, sample_args, primal - ϕqIc = intϕqI(ϕqI) - ϕqPc = intϕqP(ϕqP) - # - 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 = true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, - transP, transM) - primal = CP.neg_elbo_sites!(h0_point, CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) - dϕ, _ = CP.grad_neg_elbo_sites_enzyme( - h0_point, (;), CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix, - ) - dϕ.ϕqP -end +# @testset "grad_neg_elbo_sites_enzyme" begin +# local ϕqIc, ϕqPc, rng1, indiv_args, ϕ, sample_args, primal +# ϕqIc = intϕqI(ϕqI) +# ϕqPc = intϕqP(ϕqP) +# # +# 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 = true, pbm_covar_indices, intϕqP, intϕqI, cor_ends, +# transP, transM) +# primal = CP.neg_elbo_sites!(h0_point, CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix) +# dϕ, _ = CP.grad_neg_elbo_sites_enzyme( +# h0_point, (;), CP.NegLogDensity(), ϕ, rnormPM_point, sample_args, indiv_args, nljoint_args_fix, +# ) +# dϕ.ϕqP +# end @testset "HybridPointSolverSite" begin # in order to test inversion, set the dataloader to the synthetic data @@ -191,50 +191,43 @@ end local xM, xP, y_o, y_unc, θP_true, θMs_true, y_true (; xM, xP, y_o, y_unc, θP_true, θMs_true, y_true) = HVI.gen_hybridproblem_synthetic(rng1, prob; scenario) # constrain to training dataset - i_indiv_test = get_i_sites_test(prob; scenario) + n_indiv = size(y_o,2) + i_indiv_test = HVI.get_i_sites_test(prob; scenario) i_indiv = setdiff(1:n_indiv, i_indiv_test) (; xM, xP, y_o, y_unc, θMs_true, y_true) = map( (; xM, xP, y_o, y_unc, θMs_true, y_true)) do x - selectdim(x,ndims(x),i_indiv) + x[:,i_indiv] end - # local n_indiv, i_indiv_test, i_indiv, i_indiv_train, n_batch, train_dataloader, test_data - # n_indiv = size(xM, 2) - # i_indiv_test = HVI.get_i_sites_test(prob; scenario) - # i_indiv = setdiff(1:n_indiv, i_indiv_test) - # i_indiv_train = 1:length(i_indiv) # index into training data - # n_batch = length(h0_point.helpers_sites) - # train_dataloader = MLUtils.DataLoader(( - # CA.getdata(xM)[:,i_indiv], CA.getdata(xP)[:,i_indiv], - # y_o[:,i_indiv], y_unc[:,i_indiv], i_indiv_train); - # batchsize = n_batch, partial = false) - # test_data = (; xM=xM[:,i_indiv_test], xP=xP[:,i_indiv_test], y_o=y_o[:,i_indiv_test], y_unc=y_unc[:,i_indiv_test], itrain_sites=zeros(length(i_indiv_test))) - # prob_syn = HybridProblem(prob; scenario, train_dataloader, test_data) prob_syn = prob y_o2 = get_hybridproblem_train_dataloader(prob; scenario).data[3] - y_o .- y_o2 + @test y_o == y_o2 # make sure to use the same synthetic data for good y_true solver_point = HybridPointSolverSite(alg=Optimisers.Adam(0.0001)) - ϕ0 = (;ϕg = copy(CA.getdata(ϕg)), ϕqP=copy(ϕqP), ϕqI=copy(ϕqI)) + #ϕ0 = (;ϕg = copy(CA.getdata(ϕg)), ϕqP=copy(ϕqP), ϕqI=copy(ϕqI)) + res = CommonSolve.solve(prob_syn, solver_point, CP.NegLogDensity(), (;); scenario, + epochs=0, ); + ϕ0 = res.ϕ + ϕ0.ϕqP[1:length(θP_true)] .+= log(0.9) # start from 90% exp.(ϕ0.ϕqP[1:2]) # updated during solve - #ϕ0 = res.ϕ res = CommonSolve.solve(prob_syn, solver_point, CP.NegLogDensity(), ϕ0; scenario, - epochs=20000, + epochs=2_000, show_train_loss = Val(true) - ) - exp.(res.ϕ.ϕqP[1:2]) - #@test exp.(intϕqP(res.ϕ.ϕqP).μζP) ≈ θP_true atol = 0.1 - # run forward - ϕ_tmp = (;res.ϕ.ϕg, ϕ0.ϕqP, ϕ0.ϕqI) - hp, hi_agg = HVI.predict_point!(h0_point, prob_syn, res.ϕ; scenario) - hp0, hi_agg0 = HVI.predict_point!(h0_point, prob_syn, ϕ0; scenario) - hpt, hi_aggt = HVI.predict_point!(h0_point, prob_syn, ϕ_tmp; scenario) + ); + hp, hi_agg = HVI.predict_point!(res.h0, prob_syn, res.ϕ; scenario) @test hp.θsP[:,1] ≈ θP_true atol=0.01 + @test mean(abs.(hi_agg.θsM .- θMs_true)) < 0.01 () -> begin + exp.(res.ϕ.ϕqP[1:2]) + #@test exp.(intϕqP(res.ϕ.ϕqP).μζP) ≈ θP_true atol = 0.1 + # run forward + ϕ_tmp = (;res.ϕ.ϕg, ϕ0.ϕqP, ϕ0.ϕqI) + hp0, hi_agg0 = HVI.predict_point!(res.h0, prob_syn, ϕ0; scenario) + hpt, hi_aggt = HVI.predict_point!(res.h0, prob_syn, ϕ_tmp; scenario) # @usingany UnicodePlots i_obs = 3 #scatterplot(hi_agg.θsM[i_obs,:], θMs_true[i_obs,:]) scatterplot(y_true[i_obs,:], hi_agg.y[i_obs,:]) - scatterplot(hi_agg0.y[i_obs,:], y_true[i_obs,:]) - scatterplot(hi_aggt.y[i_obs,:], y_true[i_obs,:]) + scatterplot(y_true[i_obs,:], hi_agg0.y[i_obs,:]) + scatterplot(y_true[i_obs,:], hi_aggt.y[i_obs,:]) scatterplot(hi_aggt.y[i_obs,:], hi_agg.y[i_obs,:]) sum(abs2, hi_agg0.y .- y_true) sum(abs2, hi_agg.y .- y_true) @@ -242,10 +235,6 @@ end scatterplot(θMs_true[1,:], hi_agg.θsM[1,:], ) scatterplot(θMs_true[2,:], hi_agg.θsM[2,:]) end - @test hi_agg.θsM ≈ θMs_true atol=0.01 - h0p.helpers_sites[1].θsM - #@test res.ϕ != ϕ - x = h0p.helpers_sites[1] end