return updated problem from solvers

Author: bgctw

src/HybridSolver.jl

@@ -48,7 +48,10 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
         Optimization.AutoZygote())
     optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader)
     res = Optimization.solve(optprob, solver.alg; kwargs...)
-    (; ϕ = intϕ(res.u), resopt = res)
+    ϕ = intϕ(res.u)
+    θP = cpu_ca(apply_preserve_axes(transP, cpu_ca(ϕ).ϕP))
+    probo = update(prob; ϕg = cpu_ca(ϕ).ϕg, θP)
+    (; ϕ, resopt = res, probo)
 end
 
 struct HybridPosteriorSolver{A} <: AbstractHybridSolver
@@ -107,7 +110,9 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
     optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader)
     res = Optimization.solve(optprob, solver.alg; kwargs...)
     ϕc = interpreters.μP_ϕg_unc(res.u)
-    (; ϕ = ϕc, θP = cpu_ca(apply_preserve_axes(transP, ϕc.μP)), resopt = res, interpreters)
+    θP = cpu_ca(apply_preserve_axes(transP, ϕc.μP))
+    probo = update(prob; ϕg = cpu_ca(ϕ).ϕg, θP = θP, ϕunc = cpu_ca(ϕ).unc);
+    (; ϕ = ϕc, θP, resopt = res, interpreters, probo)
 end
 
 function fit_narrow_normal(θi, prior, θmean_quant)

src/gf.jl

@@ -29,17 +29,18 @@ function gf(prob::AbstractHybridProblem, xM, xP, args...;
     (; θP, θM) = get_hybridproblem_par_templates(prob; scenario)
     (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
     intP = ComponentArrayInterpreter(θP)
-    pbm_covar_indices = intP(1:length(intP))[pbm_covars]
+    pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario)
+    pbm_covar_indices = CA.getdata(intP(1:length(intP))[pbm_covars])
     ζP = inverse(transP)(θP)
     g_dev, ϕg_dev, ζP_dev =  (gdev(g), gdev(ϕg), gdev(CA.getdata(ζP))) 
-    gf(g_dev, transM, transP, f, xM, xP, ϕg_dev, ζP_dev; cdev, pbm_covar_indices, kwargs...)
+    gf(g_dev, transM, transP, f, xM, xP, ϕg_dev, ζP_dev, pbm_covar_indices; cdev, kwargs...)
 end
 
 function gf(g, transM, transP, f, xM, xP, ϕg, ζP; 
     cdev = identity, pbm_covars, 
-    intP = ComponentArrayInterpreter(ζP))
+    intP = ComponentArrayInterpreter(ζP), kwargs...)
     pbm_covar_indices = intP(1:length(intP))[pbm_covars]
-    gf(g, transM, transP, f, xM, xP, ϕg, ζP, pbm_covar_indices)
+    gf(g, transM, transP, f, xM, xP, ϕg, ζP, pbm_covar_indices; kwargs...)
 end
 
 
@@ -52,7 +53,6 @@ function gf(g, transM, transP, f, xM, xP, ϕg, ζP, pbm_covar_indices::AbstractV
     #     # otherwise Zyote fails on cpu_handler
     #     ζP = copy(ζP)
     # end
-    #Main.@infiltrate_main
     #xMP = _append_PBM_covars(xM, intP(ζP), pbm_covars) 
     xMP = _append_each_covars(xM, CA.getdata(ζP), pbm_covar_indices)
     θMs = gtrans(g, transM, xMP, ϕg; cdev)
@@ -99,7 +99,7 @@ function get_loss_gf(g, transM, transP, f, y_o_global,
 
     let g = g, transM = transM, transP = transP, f = f, y_o_global = y_o_global, 
         intϕ = get_concrete(intϕ),
-        pbm_covar_indices = intP(1:length(intP))[pbm_covars]
+        pbm_covar_indices = CA.getdata(intP(1:length(intP))[pbm_covars])
         #, intP = get_concrete(intP)
         #inv_transP = inverse(transP), kwargs = kwargs
         function loss_gf(p, xM, xP, y_o, y_unc, i_sites)

test/test_HybridProblem.jl

@@ -172,9 +172,9 @@ test_with_flux = (scenario) -> begin
         rng = StableRNG(111)
         solver = HybridPosteriorSolver(; alg=Adam(0.02), n_batch=11, n_MC=3)
         (; ϕ, θP, resopt) = solve(prob, solver; scenario, rng,
-            callback = callback_loss(100), maxiters = 1200,
+            #callback = callback_loss(100), maxiters = 1200,
             #maxiters = 20 # too small so that it yields error
-            #maxiters=200,
+            maxiters=37,
             θmean_quant = 0.01,   # test constraining mean to initial prediction     
             gdev = identity
         )
@@ -199,6 +199,25 @@ test_with_flux = (scenario) -> begin
             );
             @test CA.getdata(ϕ) isa GPUArraysCore.AbstractGPUVector
             @test cdev(ϕ.unc.ρsM)[1] > 0
+            #
+            solver = HybridPosteriorSolver(; alg=Adam(0.02), n_batch=11, n_MC=3)
+            test_correlation = () -> begin
+                n_epoch = 100 # requires 
+                (; ϕ, θP, resopt, probo) = solve(prob, solver; scenario = scenf,
+                    maxiters = n_batches_in_epoch * n_epoch, 
+                    callback = callback_loss(n_batches_in_epoch*5)
+                );
+                (; θ, y, entropy_ζ) = predict_gf(rng, probo; scenario = scenf, n_sample_pred = 400);            
+                mean_θ = CA.ComponentVector(mean(CA.getdata(θ); dims = 2)[:, 1], CA.getaxes(θ[:, 1])[1])
+                residθ = θ .- mean_θ
+                cr = cor(CA.getdata(residθ));
+                i_sites = [1,2,3]
+                tmp = CA.ComponentArray(collect(axes(θ[:,1],1)), CA.getaxes(θ[:,1]));
+                #ax = map(x -> axes(x,1), get_hybridproblem_par_templates(probo; scenario = scenf))
+                is = vcat(tmp.P, vec(tmp.Ms[:,i_sites]))
+                cr[is,is]
+            end
+
         end;
         # does not work with general Bijector:
         # @testset "HybridPosteriorSolver also f on gpu" begin