implement HybridPosteriorSolver

Author: bgctw

dev/doubleMM.jl

@@ -24,15 +24,16 @@ xM_cpu = xM
 if :use_Flux ∈ scenario
     xM = CuArray(xM_cpu)
 end
-get_train_loader = (rng; n_batch, kwargs...) -> MLUtils.DataLoader((xM, xP, y_o, y_unc), batchsize = n_batch)
+get_train_loader = (rng; n_batch, kwargs...) -> MLUtils.DataLoader((xM, xP, y_o, y_unc); 
+    batchsize = n_batch, partial = false)
 σ_o = exp(first(y_unc)/2)
 
 # assign the train_loader, otherwise it eatch time creates another version of synthetic data
 prob0 = HVI.update(HybridProblem(DoubleMM.DoubleMMCase(); scenario); get_train_loader)
 
 #------- pointwise hybrid model fit
-#solver = HybridPointSolver(; alg = Adam(0.02), n_batch = 30)
-solver = HybridPointSolver(; alg = Adam(0.01), n_batch = 10)
+solver = HybridPointSolver(; alg = Adam(0.02), n_batch = 30)
+#solver = HybridPointSolver(; alg = Adam(0.01), n_batch = 10)
 #solver = HybridPointSolver(; alg = Adam(), n_batch = 200)
 (; ϕ, resopt) = solve(prob0, solver; scenario,
     rng, callback = callback_loss(100), maxiters = 1200);
@@ -116,70 +117,83 @@ end
     end
 end
 
-#----------- Hybrid Variational inference 
+#----------- Hybrid Variational inference: HVI
 
 using MLUtils
 import Zygote
 
 using CUDA
 using Bijectors
 
+solver = HybridPosteriorSolver(; alg = Adam(0.01), n_batch = 60, n_MC = 3)
+#solver = HybridPointSolver(; alg = Adam(), n_batch = 200)
+(; ϕ, θP, resopt) = solve(prob0o, solver; scenario,
+    rng, callback = callback_loss(100), maxiters = 800);
+# update the problem with optimized parameters
+prob1o = HVI.update(prob0o; ϕg=cpu_ca(ϕ).ϕg, θP=θP)
+y_pred_global, y_pred, θMs = gf(prob1o, xM, xP; scenario);
+scatterplot(θMs_true[1,:], θMs[1,:])
+scatterplot(θMs_true[2,:], θMs[2,:])
+hcat(θP_true, θP) # all parameters overestimated
 
-#n_covar = get_hybridproblem_n_covar(prob; scenario)
-#, n_batch, n_θM, n_θP) = get_hybridproblem_sizes(prob; scenario)
 
-n_covar = size(xM, 1)
+() -> begin
+    #n_covar = get_hybridproblem_n_covar(prob; scenario)
+    #, n_batch, n_θM, n_θP) = get_hybridproblem_sizes(prob; scenario)
 
-#----- fit g to θMs_true
-g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
-(; transP, transM) = get_hybridproblem_transforms(prob; scenario)
+    n_covar = size(xM, 1)
 
-function loss_g(ϕg, x, g, transM)
-    ζMs = g(x, ϕg) # predict the log of the parameters
-    θMs = reduce(hcat, map(transM, eachcol(ζMs))) # transform each column
-    loss = sum(abs2, θMs .- θMs_true)
-    return loss, θMs
-end
-loss_g(ϕg0, xM, g, transM)
+    #----- fit g to θMs_true
+    g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
+    (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
 
-optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g, transM)[1],
-    Optimization.AutoZygote())
-optprob = Optimization.OptimizationProblem(optf, ϕg0);
-res = Optimization.solve(optprob, Adam(0.02), callback = callback_loss(100), maxiters = 800);
+    function loss_g(ϕg, x, g, transM)
+        ζMs = g(x, ϕg) # predict the log of the parameters
+        θMs = reduce(hcat, map(transM, eachcol(ζMs))) # transform each column
+        loss = sum(abs2, θMs .- θMs_true)
+        return loss, θMs
+    end
+    loss_g(ϕg0, xM, g, transM)
+
+    optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g, transM)[1],
+        Optimization.AutoZygote())
+    optprob = Optimization.OptimizationProblem(optf, ϕg0);
+    res = Optimization.solve(optprob, Adam(0.02), callback = callback_loss(100), maxiters = 800);
 
-ϕg_opt1 = res.u;
-l1, θMs_pred = loss_g(ϕg_opt1, xM, g, transM)
-scatterplot(vec(θMs_true), vec(θMs_pred))
+    ϕg_opt1 = res.u;
+    l1, θMs_pred = loss_g(ϕg_opt1, xM, g, transM)
+    scatterplot(vec(θMs_true), vec(θMs_pred))
 
-f = get_hybridproblem_PBmodel(prob; scenario)
-py = get_hybridproblem_neg_logden_obs(prob; scenario)
+    f = get_hybridproblem_PBmodel(prob; scenario)
+    py = get_hybridproblem_neg_logden_obs(prob; scenario)
 
-#----------- fit g and θP to y_o
-() -> begin
-    # end2end inversion
+    #----------- fit g and θP to y_o
+    () -> begin
+        # end2end inversion
 
-    int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
-        ϕg = 1:length(ϕg0), θP = par_templates.θP))
-    p = p0 = vcat(ϕg0, par_templates.θP .* 0.9)  # slightly disturb θP_true
+        int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
+            ϕg = 1:length(ϕg0), θP = par_templates.θP))
+        p = p0 = vcat(ϕg0, par_templates.θP .* 0.9)  # slightly disturb θP_true
 
-    # Pass the site-data for the batches as separate vectors wrapped in a tuple
-    train_loader = MLUtils.DataLoader((xM, xP, y_o, y_unc), batchsize = n_batch)
+        # Pass the site-data for the batches as separate vectors wrapped in a tuple
+        train_loader = MLUtils.DataLoader((xM, xP, y_o, y_unc), batchsize = n_batch)
 
-    loss_gf = get_loss_gf(g, f, y_global_o, int_ϕθP)
-    l1 = loss_gf(p0, train_loader.data...)[1]
+        loss_gf = get_loss_gf(g, f, y_global_o, int_ϕθP)
+        l1 = loss_gf(p0, train_loader.data...)[1]
 
-    optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
-        Optimization.AutoZygote())
-    optprob = OptimizationProblem(optf, p0, train_loader)
+        optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
+            Optimization.AutoZygote())
+        optprob = OptimizationProblem(optf, p0, train_loader)
 
-    res = Optimization.solve(
-        optprob, Adam(0.02), callback = callback_loss(100), maxiters = 1000)
+        res = Optimization.solve(
+            optprob, Adam(0.02), callback = callback_loss(100), maxiters = 1000)
 
-    l1, y_pred_global, y_pred, θMs = loss_gf(res.u, train_loader.data...)
-    scatterplot(vec(θMs_true), vec(θMs))
-    scatterplot(log.(vec(θMs_true)), log.(vec(θMs)))
-    scatterplot(vec(y_pred), vec(y_o))
-    hcat(par_templates.θP, int_ϕθP(res.u).θP)
+        l1, y_pred_global, y_pred, θMs = loss_gf(res.u, train_loader.data...)
+        scatterplot(vec(θMs_true), vec(θMs))
+        scatterplot(log.(vec(θMs_true)), log.(vec(θMs)))
+        scatterplot(vec(y_pred), vec(y_o))
+        hcat(par_templates.θP, int_ϕθP(res.u).θP)
+    end
 end
 
 #---------- HVI

src/AbstractHybridProblem.jl

@@ -17,7 +17,6 @@ optionally
 """
 abstract type AbstractHybridProblem end;
 
-
 """
     get_hybridproblem_MLapplicator([rng::AbstractRNG,] ::AbstractHybridProblem; scenario=())
 
@@ -28,9 +27,9 @@ returns a Tuple of
 - AbstractModelApplicator
 - initial parameter vector
 """
-function get_hybridproblem_MLapplicator end    
+function get_hybridproblem_MLapplicator end
 
-function get_hybridproblem_MLapplicator(prob::AbstractHybridProblem; scenario=())
+function get_hybridproblem_MLapplicator(prob::AbstractHybridProblem; scenario = ())
     get_hybridproblem_MLapplicator(Random.default_rng(), prob; scenario)
 end
 
@@ -56,16 +55,14 @@ function get_hybridproblem_PBmodel end
 Provide a `function(y_obs, ypred) -> Real` that computes the negative logdensity
 of the observations, given the predictions.
 """
-function get_hybridproblem_neg_logden_obs end    
-
+function get_hybridproblem_neg_logden_obs end
 
 """
     get_hybridproblem_par_templates(::AbstractHybridProblem; scenario)
 
 Provide tuple of templates of ComponentVectors `θP` and `θM`.
 """
-function get_hybridproblem_par_templates end    
-
+function get_hybridproblem_par_templates end
 
 """
     get_hybridproblem_transforms(::AbstractHybridProblem; scenario)
@@ -96,7 +93,7 @@ function get_hybridproblem_n_covar(prob::AbstractHybridProblem; scenario)
     train_loader = get_hybridproblem_train_dataloader(Random.default_rng(), prob; scenario)
     (xM, xP, y_o, y_unc) = first(train_loader)
     n_covar = size(xM, 1)
-    return(n_covar)
+    return (n_covar)
 end
 
 """
@@ -118,7 +115,7 @@ function gen_hybridcase_synthetic end
 
 Determine the FloatType for given Case and scenario, defaults to Float32
 """
-function get_hybridproblem_float_type(prob::AbstractHybridProblem; scenario=())
+function get_hybridproblem_float_type(prob::AbstractHybridProblem; scenario = ())
     return eltype(get_hybridproblem_par_templates(prob; scenario).θM)
 end
 
@@ -131,20 +128,20 @@ Return a DataLoader that provides a tuple of
 - `y_o`: matrix of observations with added noise, with one column per site
 - `y_unc`: matrix `sizeof(y_o)` of uncertainty information 
 """
-function get_hybridproblem_train_dataloader(rng::AbstractRNG, prob::AbstractHybridProblem; 
-    scenario = (), n_batch = 10)
+function get_hybridproblem_train_dataloader(rng::AbstractRNG, prob::AbstractHybridProblem;
+        scenario = (), n_batch = 10)
     (; xM, xP, y_o, y_unc) = gen_hybridcase_synthetic(rng, prob; scenario)
     xM_gpu = :use_Flux ∈ scenario ? CuArray(xM) : xM
-    train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o, y_unc), batchsize = n_batch)
-    return(train_loader)
+    train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o, y_unc);
+        batchsize = n_batch, partial = false)
+    return (train_loader)
 end
 
 function get_hybridproblem_train_dataloader(prob::AbstractHybridProblem; scenario = ())
     rng::AbstractRNG = Random.default_rng()
     get_hybridproblem_train_dataloader(rng, prob; scenario)
 end
 
-
 """
     get_hybridproblem_cor_starts(prob::AbstractHybridProblem; scenario)
 
@@ -163,8 +160,5 @@ If there is only single block of all ML-predicted parameters being correlated
 with each other then this block starts at position 1: `(P=(1,3), M=(1,))`.
 """
 function get_hybridproblem_cor_starts(prob::AbstractHybridProblem; scenario = ())
-    (P=(1,), M=(1,))
+    (P = (1,), M = (1,))
 end
-
-
-

src/HybridSolver.jl

@@ -24,12 +24,13 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
     f = get_hybridproblem_PBmodel(prob; scenario)
     y_global_o = FT[] # TODO
     loss_gf = get_loss_gf(g, transM, f, y_global_o, int_ϕθP)
+    # call loss function once
+    l1 = loss_gf(p0, first(train_loader)...)[1]
     # data1 = first(train_loader)
-    # l1 = loss_gf(p0, first(train_loader)...)[1]
     # Zygote.gradient(p0 -> loss_gf(p0, data1...)[1], p0)
     optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
         Optimization.AutoZygote())
-    optprob = OptimizationProblem(optf, p0, train_loader)
+    optprob = OptimizationProblem(optf, CA.getdata(p0), train_loader)
     res = Optimization.solve(optprob, solver.alg; kwargs...)
     (;ϕ = int_ϕθP(res.u), resopt = res)
 end
@@ -42,24 +43,54 @@ struct HybridPosteriorSolver{A} <: AbstractHybridSolver
     n_MC::Int
 
 end
-HybridPosteriorSolver(; alg, n_batch = 10, n_MC = 3) = HybridPointSolver(alg, n_batch, n_MC)
+HybridPosteriorSolver(; alg, n_batch = 10, n_MC = 3) = HybridPosteriorSolver(alg, n_batch, n_MC)
 
 function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolver; 
     scenario, rng = Random.default_rng(), kwargs...)
     par_templates = get_hybridproblem_par_templates(prob; scenario)
+    (; θP, θM) = par_templates
     g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario);
     (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
     (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
-        θP_true, θMs_true[:, 1], ϕg0, solver.n_batch; transP, transM);
+        θP, θM, ϕg0, solver.n_batch; transP, transM);
     use_gpu = (:use_Flux ∈ scenario)
-    # ϕd = use_gpu ? CuArray(ϕ) : ϕ
-    # train_loader = get_hybridproblem_train_dataloader(rng, prob; scenario, solver.n_batch)    
-    # f = get_hybridproblem_PBmodel(prob; scenario)
-    # y_global_o = Float32[] # TODO
-    # loss_gf = get_loss_gf(g, transM, f, y_global_o, int_ϕθP)
-    # optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
-    #     Optimization.AutoZygote())
-    # optprob = OptimizationProblem(optf, p0, train_loader)
-    # res = Optimization.solve(optprob, solver.alg; kwargs...)
+    ϕ0 = use_gpu ? CuArray(ϕ) : ϕ # TODO replace CuArray by something more general
+    train_loader = get_hybridproblem_train_dataloader(rng, prob; scenario, solver.n_batch)    
+    f = get_hybridproblem_PBmodel(prob; scenario)
+    py = get_hybridproblem_neg_logden_obs(prob; scenario)
+    y_global_o = Float32[] # TODO
+    loss_elbo = get_loss_elbo(g, transPMs_batch, f, py, y_global_o, interpreters; solver.n_MC)
+    # test loss function once
+    l0 = loss_elbo(ϕ0, rng, first(train_loader)...)
+    optf = Optimization.OptimizationFunction((ϕ, data) -> loss_elbo(ϕ, rng, data...)[1],
+        Optimization.AutoZygote())
+    optprob = OptimizationProblem(optf, CA.getdata(ϕ0), 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)
+end
+
+"""
+Create a loss function for parameter vector ϕ, given 
+- g(x, ϕ): machine learning model 
+- transPMS: transformation from unconstrained space to parameter space
+- f(θMs, θP): mechanistic model 
+- interpreters: assigning structure to pure vectors, see neg_elbo_transnorm_gf
+- n_MC: number of Monte-Carlo sample to approximate the expected value across distribution
+
+The loss function takes in addition to ϕ, data that changes with minibatch
+- rng: random generator
+- xM: matrix of covariates, sites in columns
+- xP: drivers for the processmodel: Iterator of size n_site
+- y_o, y_unc: matrix of observations and uncertainties, sites in columns
+"""
+function get_loss_elbo(g, transPMs, f, py, y_o_global, interpreters; n_MC)
+    let g = g, transPMs = transPMs, f = f, py=py, y_o_global = y_o_global, n_MC = n_MC
+        interpreters = map(get_concrete, interpreters)
+        function loss_elbo(ϕ, rng, xM, xP, y_o, y_unc)
+            neg_elbo_transnorm_gf(rng, ϕ, g, transPMs, f, py,
+            xM, xP, y_o, y_unc, interpreters; n_MC)
+        end
+    end
 end
 

src/elbo.jl

@@ -87,6 +87,8 @@ function generate_ζ(rng, g, ϕ::AbstractVector, xM::AbstractMatrix,
     μ_ζMs0 = g(xM, ϕg) # TODO provide μ_ζP to g
     ζ_resid, σ = sample_ζ_norm0(rng, μ_ζP, μ_ζMs0, ϕc.unc; n_MC, cor_starts)
     #ζ_resid, σ = sample_ζ_norm0(rng, ϕ[1:2], reshape(ϕ[2 .+ (1:20)],2,:), ϕ[(end-length(interpreters.unc)+1):end], interpreters.unc; n_MC)
+    # @show size(ζ_resid)
+    # @show length(interpreters.PMs)
     ζ = stack(map(eachcol(ζ_resid)) do r
         rc = interpreters.PMs(r)
         ζP = μ_ζP .+ rc.P

src/gf.jl

@@ -52,7 +52,7 @@ Create a loss function for parameter vector p, given
 - int_ϕθP: interpreter attachin axis with compponents ϕg and pc.θP
 """
 function get_loss_gf(g, transM, f, y_o_global, int_ϕθP::AbstractComponentArrayInterpreter)
-    let g = g, transM = transM, f = f, int_ϕθP = int_ϕθP
+    let g = g, transM = transM, f = f, int_ϕθP = int_ϕθP, y_o_global = y_o_global
         function loss_gf(p, xM, xP, y_o, y_unc)
             σ = exp.(y_unc ./ 2)
             pc = int_ϕθP(p)

src/init_hybrid_params.jl

@@ -37,7 +37,7 @@ function init_hybrid_params(θP, θM, ϕg, n_batch;
         ρsP,
         ρsM)
     ϕ = CA.ComponentVector(;
-        μP = inverse(transP)(θP),
+        μP = apply_preserve_axes(inverse(transP),θP),
         ϕg = ϕg,
         unc = ϕunc0);
     #

src/util_ca.jl

@@ -6,4 +6,8 @@ Move ComponentArray form gpu to cpu.
 function cpu_ca end
 # define in FluxExt
 
+function apply_preserve_axes(f, ca::CA.ComponentArray)
+    CA.ComponentArray(f(ca), CA.getaxes(ca))
+end
+