test point fit and devide logJacTM by n_indiv

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

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) 
@@ -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,9 +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);
-    rnormPM = prepare_rnorm(ϕqP; n_θP, n_θM, n_MC=NMC, n_site)
+    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, ϕ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(ϕ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)
@@ -85,7 +96,7 @@ 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)
+    ϕ = 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 +105,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 +116,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 +137,100 @@ 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
+    (;ϕ, h0, h0_test)
+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
 

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,14 +287,29 @@ 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
         res0 .* app.width .+ app.offset
     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)

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)

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: 
@@ -93,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
@@ -107,7 +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]
-    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)