test distribution oof generated residuals

Author: bgctw

dev/doubleMM.jl

@@ -160,8 +160,8 @@ n_epoch = 40
 # update the problem with optimized parameters, including uncertainty
 prob1o = probo;
 n_sample_pred = 400
-#(; θ, y) = predict_gf(rng, prob1o, xM, xP; scenario, n_sample_pred);
-(; θ, y) = predict_gf(rng, prob1o; scenario, n_sample_pred);
+#(; θ, y) = predict_hvi(rng, prob1o, xM, xP; scenario, n_sample_pred);
+(; θ, y) = predict_hvi(rng, prob1o; scenario, n_sample_pred);
 (θ1, y1) = (θ, y);
 
 () -> begin # prediction with fitted parameters  (should be smaller than mean)
@@ -210,7 +210,7 @@ end
     prob2o_indep = HVI.update(tmp["prob2o"], get_train_loader = prob0.get_train_loader);
     # test predicting correct obs-uncertainty of predictive posterior
     n_sample_pred = 400
-    (; θ, y, entropy_ζ) = predict_gf(rng, prob2o_indep, xM, xP; scenario, n_sample_pred);
+    (; θ, y, entropy_ζ) = predict_hvi(rng, prob2o_indep, xM, xP; scenario, n_sample_pred);
     (θ2_indep, y2_indep) = (θ, y)
     #(θ2_indep, y2_indep) = (θ2, y2)  # workaround to use covarK2 when loading failed
 end
@@ -241,12 +241,12 @@ end
 #ζMs_VI = g_flux(xM_gpu, ζ_VIc.ϕg |> Flux.gpu) |> Flux.cpu
 ϕunc_VI = interpreters.unc(ζ_VIc.unc)
 ϕunc_VI.ρsM
-exp.(ϕunc_VI.logσ2_logP)
+exp.(ϕunc_VI.logσ2_ζP)
 exp.(ϕunc_VI.coef_logσ2_ζMs[1, :])
 
 # test predicting correct obs-uncertainty of predictive posterior
 n_sample_pred = 400
-(; θ, y, entropy_ζ) = predict_gf(rng, prob2o; scenario, n_sample_pred);
+(; θ, y, entropy_ζ) = predict_hvi(rng, prob2o; scenario, n_sample_pred);
 (θ2, y2) = (θ, y)
 size(y) # n_obs x n_site, n_sample_pred
 size(θ)  # n_θP + n_site * n_θM x n_sample
@@ -320,7 +320,7 @@ end
 
 () -> begin # look at distribution of parameters, predictions, and likelihood and elob at one site
     function predict_site(probo, i_site)
-        (; θ, y, entropy_ζ) = predict_gf(rng, probo, xM, xP; scenario, n_sample_pred)
+        (; θ, y, entropy_ζ) = predict_hvi(rng, probo, xM, xP; scenario, n_sample_pred)
         y_site = y[:, i_site, :]
         θMs_i = map(i_rep -> θ[:Ms, i_rep][:, i_site], axes(θ, 2))
         r1s = map(x -> x[1], θMs_i)

src/HybridVariationalInference.jl

@@ -84,7 +84,7 @@ include("logden_normal.jl")
 export get_ca_starts, get_ca_ends, get_cor_count
 include("cholesky.jl")
 
-export neg_elbo_gtf, predict_gf
+export neg_elbo_gtf, predict_hvi
 include("elbo.jl")
 
 export init_hybrid_params, init_hybrid_ϕunc

src/elbo.jl

@@ -141,21 +141,40 @@ end
 end
 
 """
-    predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, interpreters;
+    predict_hvi([rng], prob::AbstractHybridProblem [,xM, xP]; scenario, ...)
+    predict_hvi(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix;
         get_transPMs, get_ca_int_PMs, n_sample_pred=200, gdev = identity)
 
-Prediction function for hybrid model. Returns an NamedTuple with entries
-- `θ`: ComponentArray `(n_θP + n_site * n_θM), n_sample_pred)` of PBM model parameters.
+Prediction function for hybrid variational inference parameter model. 
+
+## Arguments
+- The problem for which to predict
+- xM: covariates for the machine-learning model (ML): Matrix (n_θM x n_site_pred).
+- xP: model drivers for process based model (PBM): Matrix with (n_site_pred) rows.
+  If provided a ComponentArray with a Tuple-Axis in rows, the PBM model can
+  access parts of it, e.g. `xP[:S1,...]`.
+
+## Keyword arguments
+- scenario
+- n_sample_pred
+
+Returns an NamedTuple `(; y, θsP, θsMs, entropy_ζ)` with entries
 - `y`: Array `(n_obs, n_site, n_sample_pred)` of model predictions.
+- `θsP`: ComponentArray `(n_θP, n_sample_pred)` of PBM model parameters
+  that are kept constant across sites.
+- `θsMs`: ComponentArray `(n_site, n_θM, n_sample_pred)` of PBM model parameters
+  that vary by site.
+- `entropy_ζ`: The entroy of the log-determinant of the transformation of 
+  the set of model parameters, which is involved in uncertainty quantification.
 """
-function predict_gf(rng, prob::AbstractHybridProblem; scenario, kwargs...)
+function predict_hvi(rng, prob::AbstractHybridProblem; scenario, kwargs...)
     dl = get_hybridproblem_train_dataloader(prob; scenario)
     dl_dev = gdev_hybridproblem_dataloader(dl; scenario)
     # predict for all sites
     xM, xP = dl_dev.data[1:2]
-    predict_gf(rng, prob, xM, xP; scenario, kwargs...)
+    predict_hvi(rng, prob, xM, xP; scenario, kwargs...)
 end
-function predict_gf(rng, prob::AbstractHybridProblem, xM::AbstractMatrix, xP;
+function predict_hvi(rng, prob::AbstractHybridProblem, xM::AbstractMatrix, xP;
     scenario,
     n_sample_pred=200,
     gdev=:use_gpu ∈ _val_value(scenario) ? gpu_device() : identity,
@@ -184,12 +203,12 @@ function predict_gf(rng, prob::AbstractHybridProblem, xM::AbstractMatrix, xP;
     int_unc = interpreters.unc
     transMs = StackedArray(transM, n_batch)        
     g_dev, ϕ_dev = gdev(g), gdev(ϕ)
-    predict_gf(rng, g_dev, f, ϕ_dev, xM, xP;
+    predict_hvi(rng, g_dev, f, ϕ_dev, xM, xP;
         int_μP_ϕg_unc, int_unc, transP, transM, 
         n_sample_pred, cdev, cor_ends, pbm_covar_indices, kwargs...)
 end
 
-function predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, xP;
+function predict_hvi(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, xP;
     int_μP_ϕg_unc::AbstractComponentArrayInterpreter,
     int_unc::AbstractComponentArrayInterpreter,
     transP, transM,
@@ -274,6 +293,8 @@ function generate_ζ(rng, g, ϕ::AbstractVector{FT}, xM::MT;
     # first pass: append μ_ζP_to covars, need ML prediction for magnitude of ζMs
     # TODO replace pbm_covar_indices by ComponentArray? dimensions to be type-inferred?
     xMP0 = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices)
+    #Main.@infiltrate_main
+
     μ_ζMs0 = g(xMP0, ϕg)::MT # for gpu restructure returns Any, so apply type
     ζP_resids, ζMs_parfirst_resids, σ = sample_ζresid_norm(rng, μ_ζP, μ_ζMs0, ϕc.unc; n_MC, cor_ends, int_unc)
     if pbm_covar_indices isa SA.SVector{0}
@@ -362,7 +383,7 @@ ML-model predcitions of size `(n_θM, n_site)`.
 
 ## Arguments
 * `int_unc`: Interpret vector as ComponentVector with components
-   ρsP, ρsM, logσ2_logP, coef_logσ2_ζMs(intercept + slope), 
+   ρsP, ρsM, logσ2_ζP, coef_logσ2_ζMs(intercept + slope), 
 """
 function sample_ζresid_norm(rng::Random.AbstractRNG, ζP::AbstractVector, ζMs::AbstractMatrix,
     args...; n_MC, cor_ends, int_unc)
@@ -392,10 +413,10 @@ function sample_ζresid_norm(urandn::AbstractMatrix, ζP::TP, ζMs::TM,
     UM = transformU_block_cholesky1(ρsM, cor_ends.M)
     cf = ϕuncc.coef_logσ2_ζMs
     logσ2_logMs = vec(cf[1, :] .+ cf[2, :] .* ζMs)
-    logσ2_logP = vec(CA.getdata(ϕuncc.logσ2_logP))
+    logσ2_ζP = vec(CA.getdata(ϕuncc.logσ2_ζP))
     # CUDA cannot multiply BlockDiagonal * Diagonal, construct already those blocks
     σMs = reshape(exp.(logσ2_logMs ./ 2), n_θM, :)
-    σP = exp.(logσ2_logP ./ 2)
+    σP = exp.(logσ2_ζP ./ 2)
     # BlockDiagonal does work with CUDA, but not with combination of Zygote and CUDA
     # need to construct full matrix for CUDA
     Uσ = _create_blockdiag(UP, UM, σP, σMs, n_batch)

src/init_hybrid_params.jl

@@ -72,7 +72,7 @@ function init_hybrid_params(θP::AbstractVector{FT}, θM::AbstractVector{FT},
 end
 
 """
-    init_hybrid_ϕunc(cor_ends, ρ0=0f0; logσ2_logP, coef_logσ2_ζMs, ρsP, ρsM)
+    init_hybrid_ϕunc(cor_ends, ρ0=0f0; logσ2_ζP, coef_logσ2_ζMs, ρsP, ρsM)
 
 Initialize vector of additional parameter of the approximate posterior.
 
@@ -83,7 +83,7 @@ Arguments:
 - `coef_logσ2_logM`: default column for `coef_logσ2_ζMs`, defaults to `[-10.0, 0.0]`
 
 Returns a `ComponentVector` of 
-- `logσ2_logP`: vector of log-variances of ζP (on log scale).
+- `logσ2_ζP`: vector of log-variances of ζP (on log scale).
   defaults to -10
 - `coef_logσ2_ζMs`: offset and slope for the log-variances of ζM scaling with 
    its value given by columns for each parameter in ζM, defaults to `[-10, 0]`
@@ -94,14 +94,14 @@ function init_hybrid_ϕunc(
         cor_ends::NamedTuple,
         ρ0::FT = 0.0f0,
         coef_logσ2_logM::AbstractVector{FT} = FT[-10.0, 0.0];
-        logσ2_logP::AbstractVector{FT} = fill(FT(-10.0), cor_ends.P[end]),
+        logσ2_ζP::AbstractVector{FT} = fill(FT(-10.0), cor_ends.P[end]),
         coef_logσ2_ζMs::AbstractMatrix{FT} = reduce(
             hcat, (coef_logσ2_logM for _ in 1:cor_ends.M[end])),
         ρsP = fill(ρ0, get_cor_count(cor_ends.P)),
         ρsM = fill(ρ0, get_cor_count(cor_ends.M)),
 ) where {FT}
     nt = (;
-        logσ2_logP,
+        logσ2_ζP,
         coef_logσ2_ζMs,
         ρsP,
         ρsM)

test/test_HybridProblem.jl

@@ -263,7 +263,7 @@ test_with_flux = (scenario) -> begin
             @test cdev(ϕ.unc.ρsM)[1] > 0 
             @test probo.ϕunc == cdev(ϕ.unc)
             n_sample_pred = 22
-            (; y, θsP, θsMs) = predict_gf(
+            (; y, θsP, θsMs) = predict_hvi(
                 rng, probo; scenario = scenf, n_sample_pred, is_inferred=Val(true));            
             (_xM, _xP, _y_o, _y_unc, _i_sites) = get_hybridproblem_train_dataloader(prob; scenario).data
             @test size(y) == (size(_y_o)..., n_sample_pred)
@@ -279,7 +279,7 @@ test_with_flux = (scenario) -> begin
                 @test probo.ϕunc == cdev(ϕ.unc)
                 # predict using problem and its associated dataloader
                 n_sample_pred = 201
-                (; y, θsP, θsMs) = predict_gf(rng, probo; scenario = scenf, n_sample_pred);            
+                (; y, θsP, θsMs) = predict_hvi(rng, probo; scenario = scenf, n_sample_pred);            
                 # to inspect correlations among θP and θMs construct ComponentVector
                 hpints = HybridProblemInterpreters(prob; scenario)
                 int_mPMs = stack_ca_int(Val((n_sample_pred,)), get_int_PMst_site(hpints))
@@ -317,7 +317,7 @@ test_with_flux = (scenario) -> begin
             );
             @test CA.getdata(ϕ) isa GPUArraysCore.AbstractGPUVector
             n_sample_pred = 11
-            (; y, θsP, θsMs) = predict_gf(
+            (; y, θsP, θsMs) = predict_hvi(
                 rng, probo; scenario = scenf, n_sample_pred,is_inferred = Val(true));
             # @test cdev(ϕ.unc.ρsM)[1] > 0 # too few iterations
         end;