Inprogress

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

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)
 

src/ModelApplicator.jl

@@ -177,7 +177,7 @@
 """
 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 @@
     # @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 @@
     # 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)

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
@@ -23,50 +23,41 @@
    of parameters into fixed, global, and individual.
     - `θc`: parameters: `ComponentMatrix` (n_par x n_MC) with each row a parameter vector
     - `xPc`: observations: `ComponentVector (n_obs) with observationsfor one site
 - `n_MC`: number of parameterizations, i.e. colums in `θsM`
 - `θP`: `ComponentVector` template of global process model parameters
 - `θM`: `ComponentVector` template of individual process model parameters
 - `θFix`: `ComponentVector` of actual fixed process model parameters
 - `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 @@
     #     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
 

src/elbo_site.jl

@@ -44,26 +44,28 @@
     ϕ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
     #res_site = map(compute_nelboi_z!, h.helpers_sites, rnormPM.M, i_sites_train, ϕm_it)
     #MAYBE: distributed mapreduce: 
     #   https://docs.julialang.org/en/v1/stdlib/Distributed/#Distributed.@distributed
     #   https://github.com/SupaeroDataScience/DE/blob/main/notebooks/Introduction%20to%20MapReduce.ipynb
-    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,10 +74,9 @@
     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
         template = eltype(ϕqIc) <: ForwardDiff.Dual ? ϕqIc : θsP
         hi = map_leaves_nt(x -> PAT.get_tmp(x, template), hi)
     end
@@ -84,7 +85,7 @@
     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) 
@@ -307,13 +308,12 @@
     (;xP, y_o, y_unc, i_indiv_train) = nljoint_args_indi
     @assert size(y_o) == size(y_unc) == (n_obs,)
     @assert size(θsP,2) == size(θsM,2) == n_MC
     (;f!,fLogLik) = nljoint_args_fix
     #θsPc = CA.ComponentMatrix(θsP, axθP, CA.FlatAxis())
     #θsMc = CA.ComponentMatrix(θsM, axθM, CA.FlatAxis())
     apply_model!(y_pred, f!, θsP, θsM, xP)
     nL_reg = sum(y_pred) / n_MC
-    nL = fLogLik(y_pred, y_o, y_unc)
-
+    nL = fLogLik(y_pred, y_o, y_unc) / n_MC
     # nL = (5 * sum(θsP) + 3 * sum(θsM)) / n_MC
     (; nL,)
     # loss_comps = neg_elbo_ζtf(
@@ -363,8 +363,7 @@
     pred
 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
 

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)