diff --git a/src/DoubleMM/f_doubleMM.jl b/src/DoubleMM/f_doubleMM.jl index 82cadc2..afd8b7f 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 @@ -226,10 +237,13 @@ 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 + nothing end g = if (:use_rangescaling ∈ scen) RangeScalingModelApplicator(g_nomag, lowers, uppers, eltype(ϕ_g0); range_scaled) @@ -292,6 +306,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 = 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) PBMPopulationApplicatorTr(f_doubleMM_sites, n_batch; pt.θP, pt.θM, θFix, xPvec) @@ -324,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) @@ -410,7 +431,7 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase; #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) + (y_true, _) = f(θP, θMs_true', xP) σ_o = FloatType(0.01) #σ_o = FloatType(0.002) logσ2_o = FloatType(2) .* log.(σ_o) @@ -471,6 +492,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..8a78539 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,21 @@ 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 +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/ModelApplicator.jl b/src/ModelApplicator.jl index b6f7d12..6a9a95e 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 @@ -231,8 +231,43 @@ function apply_model(app::NormalScalingModelApplicator, x, ϕ; kwargs...) # end end -index_firstdim(v::AbstractVector, i) = v[i] -index_firstdim(v::AbstractMatrix, i) = v[i,:] +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 !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, 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 + y +end + + +@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 df3394a..e8b1ab4 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,15 +77,47 @@ 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) <: 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 +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(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{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(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) + 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 7dde816..a9a07be 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) + 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, ϕqIc, θsP, cor_ends.M, transM, nljoint_args_fix) end @@ -55,15 +56,16 @@ 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) #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) + @@ -72,7 +74,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 @@ -84,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) @@ -312,8 +313,7 @@ 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,) # loss_comps = neg_elbo_ζtf( @@ -364,7 +364,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/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/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..679bdfc 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,17 @@ function neg_logden_indep_normal(obs::AbstractArray{TO,D}, μ::AbstractArray{TP, end +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) + (σ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 new file mode 100644 index 0000000..1ba949b --- /dev/null +++ b/test/test_elbo_site_doubleMM.jl @@ -0,0 +1,250 @@ +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 StaticArrays as SA + +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) +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) + ϕ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) + 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), + 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)) + # + 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], + i_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(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) + +#------------------------------------------------------------- + +# @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 = true + # + 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, 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) + function loop_grad_neg_elbo_sites( + n, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + for _ in 1:n + @noinline CP.grad_neg_elbo_sites( + h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + end + nothing + end + function alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + # avoid global variables -> pass them through function + loop_grad_neg_elbo_sites(1, h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + #@profview_allocs loop_grad_neg_elbo_sites(1_000,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + @test (@allocated loop_grad_neg_elbo_sites(100,h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val)) <= 6_748_800 + end + alloc_grad_neg_elbo_sites(h0, gradh0, ϕ, rnormPM, sample_args, indiv_args, nljoint_args_fix; n_MC_val) + #_i_sites_train = 1:n_site + #@usingany BenchmarkTools + #@benchmark CP.grad_neg_elbo_sites($h2, $gradh2, $rnormPM, $ϕg2v, $ϕqP, $ϕqI, $g2, $pbm_covar_indices2; i_sites_train = _i_sites_train, intϕqP=$intϕqP, intϕqI=$intϕqI, xM=$xM, is_testmode = false) + +end + + + + + +