encapsulate generating transformations and interpreters

implement predict_gf function

Author: bgctw

dev/doubleMM.jl

@@ -53,8 +53,8 @@ scatterplot(vec(θMs_true), vec(loss_g(ϕg_opt1, xM, g)[2]))
 
 f = gen_hybridcase_PBmodel(case; scenario)
 
+#----------- fit g and θP to y_o
 () -> begin
-    #----------- fit g and θP to y_o
     # end2end inversion
 
     int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
@@ -82,67 +82,78 @@ f = gen_hybridcase_PBmodel(case; scenario)
 end
 
 #---------- HVI
-# TODO think about good general initializations
-coef_logσ2_logMs = [-5.769 -3.501; -0.01791 0.007951]
-logσ2_logP = CA.ComponentVector(r0 = -8.997, K2 = -5.893)
-mean_σ_o_MC = 0.006042
-
-# correlation matrices
-ρsP = zeros(sum(1:(n_θP - 1)))
-ρsM = zeros(sum(1:(n_θM - 1)))
-
-ϕunc = CA.ComponentVector(;
-    logσ2_logP = logσ2_logP,
-    coef_logσ2_logMs = coef_logσ2_logMs,
-    ρsP,
-    ρsM)
-int_unc = ComponentArrayInterpreter(ϕunc)
-
-# for a conservative uncertainty assume σ2=1e-10 and no relationship with magnitude
-ϕunc0 = CA.ComponentVector(;
-    logσ2_logP = fill(-10.0, n_θP),
-    coef_logσ2_logMs = reduce(hcat, ([-10.0, 0.0] for _ in 1:n_θM)),
-    ρsP,
-    ρsM)
-
 logσ2y = 2 .* log.(σ_o)
 n_MC = 3
+(; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
+    θP_true, θMs_true[:, 1], ϕg_opt1, n_batch; transP = asℝ₊, transM = asℝ₊);
+ϕ_true = ϕ
 
-transPMs_batch = as(
-    (P = as(Array, asℝ₊, n_θP),
-    Ms = as(Array, asℝ₊, n_θM, n_batch)))
-transPMs_all = as(
-    (P = as(Array, asℝ₊, n_θP),
-    Ms = as(Array, asℝ₊, n_θM, n_site)))
-
-n_ϕg = length(ϕg_opt1)
-ϕt_true = θ = CA.ComponentVector(;
-    μP = θP_true,
-    ϕg = ϕg_opt1,
-    unc = ϕunc);
-trans_gu = as(
-    (μP = as(Array, asℝ₊, n_θP),
-    ϕg = as(Array, n_ϕg),
-    unc = as(Array, length(ϕunc))))
-trans_g = as(
-    (μP = as(Array, asℝ₊, n_θP),
-    ϕg = as(Array, n_ϕg)))
-
-#const 
-int_PMs_batch = ComponentArrayInterpreter(CA.ComponentVector(; θP = θP_true,
-    θMs = CA.ComponentMatrix(
-        zeros(n_θM, n_batch), first(CA.getaxes(θMs_true)), CA.Axis(i = 1:n_batch))))
-
-interpreters = interpreters_g = map(get_concrete,
-    (;
-        μP_ϕg_unc = ComponentArrayInterpreter(ϕt_true),
-        PMs = int_PMs_batch,
-        unc = ComponentArrayInterpreter(ϕunc)
-    ))
-
-ϕ_true = inverse_ca(trans_gu, ϕt_true)
+() -> begin
+    coef_logσ2_logMs = [-5.769 -3.501; -0.01791 0.007951]
+    logσ2_logP = CA.ComponentVector(r0 = -8.997, K2 = -5.893)
+    mean_σ_o_MC = 0.006042
+
+    # correlation matrices
+    ρsP = zeros(sum(1:(n_θP - 1)))
+    ρsM = zeros(sum(1:(n_θM - 1)))
+
+    ϕunc = CA.ComponentVector(;
+        logσ2_logP = logσ2_logP,
+        coef_logσ2_logMs = coef_logσ2_logMs,
+        ρsP,
+        ρsM)
+    int_unc = ComponentArrayInterpreter(ϕunc)
+
+    # for a conservative uncertainty assume σ2=1e-10 and no relationship with magnitude
+    ϕunc0 = CA.ComponentVector(;
+        logσ2_logP = fill(-10.0, n_θP),
+        coef_logσ2_logMs = reduce(hcat, ([-10.0, 0.0] for _ in 1:n_θM)),
+        ρsP,
+        ρsM)
+
+    transPMs_batch = as(
+        (P = as(Array, asℝ₊, n_θP),
+        Ms = as(Array, asℝ₊, n_θM, n_batch)))
+    transPMs_allsites = as(
+        (P = as(Array, asℝ₊, n_θP),
+        Ms = as(Array, asℝ₊, n_θM, n_site)))
+
+    n_ϕg = length(ϕg_opt1)
+    ϕt_true = θ = CA.ComponentVector(;
+        μP = θP_true,
+        ϕg = ϕg_opt1,
+        unc = ϕunc)
+    trans_gu = as(
+        (μP = as(Array, asℝ₊, n_θP),
+        ϕg = as(Array, n_ϕg),
+        unc = as(Array, length(ϕunc))))
+    trans_g = as(
+        (μP = as(Array, asℝ₊, n_θP),
+        ϕg = as(Array, n_ϕg)))
+
+    #const 
+    int_PMs_batch = ComponentArrayInterpreter(CA.ComponentVector(; θP = θP_true,
+        θMs = CA.ComponentMatrix(
+            zeros(n_θM, n_batch), first(CA.getaxes(θMs_true)), CA.Axis(i = 1:n_batch))))
+
+    interpreters = interpreters_g = map(get_concrete,
+        (;
+            μP_ϕg_unc = ComponentArrayInterpreter(ϕt_true),
+            PMs = int_PMs_batch,
+            unc = ComponentArrayInterpreter(ϕunc)
+        ))
+
+    ϕ_true = inverse_ca(trans_gu, ϕt_true)
+end
+
+ϕ_ini0 = ζ = vcat(ϕ_true[:μP] .* 0.0, ϕg0, ϕ_true[[:unc]]); # scratch
+#
+# true values
 ϕ_ini = ζ = vcat(ϕ_true[[:μP, :ϕg]] .* 1.2, ϕ_true[[:unc]]); # slight disturbance
-ϕ_ini0 = ζ = vcat(ϕ_true[:μP] .* 0.0, ϕg0, ϕunc0); # scratch
+# hardcoded from HMC inversion
+ϕ_ini.unc.coef_logσ2_logMs = [-5.769 -3.501; -0.01791 0.007951]
+ϕ_ini.unc.logσ2_logP = CA.ComponentVector(r0 = -8.997, K2 = -5.893)
+mean_σ_o_MC = 0.006042
 
 # test cost function and gradient
 () -> begin
@@ -161,10 +172,10 @@ end
     train_loader = MLUtils.DataLoader((xM, y_o), batchsize = n_batch)
 
     optf = Optimization.OptimizationFunction(
-        (ζg, data) -> begin
+        (ϕ, data) -> begin
             xM, y_o = data
             neg_elbo_transnorm_gf(
-                rng, g, f, ζg, y_o, xM, transPMs_batch,
+                rng, g, f, ϕ, y_o, xM, transPMs_batch,
                 map(get_concrete, interpreters_g); n_MC = 5, logσ2y)
         end,
         Optimization.AutoZygote())
@@ -181,7 +192,7 @@ g_flux, ϕg0_flux_cpu = gen_hybridcase_MLapplicator(case, FluxMLengine; scenario
 
 # otpimize using LUX
 () -> begin
-    using Lux
+    #using Lux
     g_lux = Lux.Chain(
         # dense layer with bias that maps to 8 outputs and applies `tanh` activation
         Lux.Dense(n_covar => n_covar * 4, tanh),
@@ -208,18 +219,19 @@ function fcost(ϕ)
         n_MC = 8, logσ2y = logσ2y)
 end
 fcost(ϕ)
-Zygote.gradient(fcost, ϕ) |> cpu;
+#Zygote.gradient(fcost, ϕ) |> cpu;
 gr = Zygote.gradient(fcost, CA.getdata(ϕ));
-gr_c = CA.ComponentArray(gr[1], CA.getaxes(ϕ)...)
+gr_c = CA.ComponentArray(gr[1] |> Flux.cpu, CA.getaxes(ϕ)...)
 
 train_loader = MLUtils.DataLoader((xM_gpu, y_o), batchsize = n_batch)
 
 optf = Optimization.OptimizationFunction(
-    (ζg, data) -> begin
+    (ϕ, data) -> begin
         xM, y_o = data
-        neg_elbo_transnorm_gf(
-            rng, g_flux, f, ζg, y_o, xM, transPMs_batch,
-            map(get_concrete, interpreters_g); n_MC = 5, logσ2y)
+        fcost(ϕ)
+        # neg_elbo_transnorm_gf(
+        #     rng, g_flux, f, ϕ, y_o, xM, transPMs_batch,
+        #     map(get_concrete, interpreters); n_MC = 5, logσ2y)
     end,
     Optimization.AutoZygote())
 optprob = Optimization.OptimizationProblem(
@@ -230,40 +242,31 @@ res = res_gpu = Optimization.solve(
 # start from zero 
 () -> begin
     optprob = Optimization.OptimizationProblem(
-        optf, CA.getdata(ϕ_ini0) |> Flux.gpu, train_loader);
+        optf, CA.getdata(ϕ_ini0) |> Flux.gpu, train_loader)
     res = res_gpu = Optimization.solve(
-            optprob, Optimisers.Adam(0.02), callback = callback_loss(50), maxiters = 4_000);
+        optprob, Optimisers.Adam(0.02), callback = callback_loss(50), maxiters = 4_000)
 end
 
-ζ_VIc = interpreters_g.μP_ϕg_unc(res.u |> Flux.cpu)
-ζMs_VI = g(xM, ζ_VIc.ϕg)
-ϕunc_VI = int_unc(ζ_VIc.unc)
+ζ_VIc = interpreters.μP_ϕg_unc(res.u |> Flux.cpu)
+ζMs_VI = g_flux(xM_gpu, ζ_VIc.ϕg |> Flux.gpu) |> Flux.cpu
+ϕunc_VI = interpreters.unc(ζ_VIc.unc)
 
 hcat(θP_true, exp.(ζ_VIc.μP))
 plt = scatterplot(vec(θMs_true), vec(exp.(ζMs_VI)))
 #lineplot!(plt, 0.0, 1.1, identity)
 # 
-hcat(ϕunc, ϕunc_VI) # need to compare to MC sample
+hcat(ϕ_ini.unc, ϕunc_VI) # need to compare to MC sample
 # hard to estimate for original very small theta's but otherwise good
 
 # test predicting correct obs-uncertainty of predictive posterior
 # TODO reuse g_flux rather than g
 n_sample_pred = 200
-intm_PMs_gen = ComponentArrayInterpreter(CA.ComponentVector(; θP = θP_true,
-    θMs = CA.ComponentMatrix(
-        zeros(n_θM, n_site), first(CA.getaxes(θMs_true)), CA.Axis(i = 1:n_sample_pred))))
-
-ζs, _ = HVI.generate_ζ(rng, g, f, res.u |> Flux.cpu, xM,
-    (; interpreters..., PMs = intm_PMs_gen); n_MC = n_sample_pred)
-# ζ = ζs[:,1]   
-θsc = stack(
-    ζ -> CA.getdata(CA.ComponentVector(
-        TransformVariables.transform(transPMs_all, ζ))),
-    eachcol(ζs));
-y_pred = stack(map(ζ -> first(HVI.predict_y(ζ, f, transPMs_all)), eachcol(ζs)));
-
-size(y_pred)
-σ_o_post = mapslices(std, y_pred; dims = 3)[:, :, 1];
+y_pred = predict_gf(rng, g_flux, f, res.u, xM_gpu, interpreters;
+    get_transPMs, get_ca_int_PMs, n_sample_pred);
+size(y_pred) # n_obs x n_site, n_sample_pred
+
+σ_o_post = dropdims(std(y_pred; dims = 3), dims=3)
+
 #describe(σ_o_post)
 hcat(σ_o, fill(mean_σ_o_MC, length(σ_o)),
     mean(σ_o_post, dims = 2), sqrt.(mean(abs2, σ_o_post, dims = 2)))

src/HybridVariationalInference.jl

@@ -44,9 +44,12 @@ include("logden_normal.jl")
 #export - all internal
 include("cholesky.jl")
 
-export neg_elbo_transnorm_gf
+export neg_elbo_transnorm_gf, predict_gf
 include("elbo.jl")
 
+export init_hybrid_params
+include("init_hybrid_params.jl")
+
 export DoubleMM
 include("DoubleMM/DoubleMM.jl")
 

src/elbo.jl

@@ -21,34 +21,37 @@ expected value of the likelihood of observations.
 function neg_elbo_transnorm_gf(rng, g, f, ϕ::AbstractVector, y_ob, x::AbstractMatrix,
     transPMs, interpreters::NamedTuple; 
     n_MC=3, logσ2y, gpu_data_handler = get_default_GPUHandler())
-    ζ, logdetΣ = generate_ζ(rng, g, f, ϕ, x, interpreters; n_MC)
-    ζ_cpu = gpu_data_handler(ζ) # differentiable fetch to CPU in Flux package extension
-    #ζi = first(eachcol(ζ_cpu))
-    nLy = reduce(+, map(eachcol(ζ_cpu)) do ζi
+    ζs, logdetΣ = generate_ζ(rng, g, f, ϕ, x, interpreters; n_MC)
+    ζs_cpu = gpu_data_handler(ζs) # differentiable fetch to CPU in Flux package extension
+    #ζi = first(eachcol(ζs_cpu))
+    nLy = reduce(+, map(eachcol(ζs_cpu)) do ζi
         y_pred_i, logjac = predict_y(ζi, f, transPMs)
         nLy1 = neg_logden_indep_normal(y_ob, y_pred_i, logσ2y)
         nLy1 - logjac
     end) / n_MC
-    ent = entropy_MvNormal(size(ζ, 1), logdetΣ)  # defined in logden_normal
-    nLy - ent
-end
-
-function predict_gf(rng, g, f, ϕ::AbstractVector, x::AbstractMatrix,
-    transPMs, interpreters::NamedTuple; 
-    n_MC=3, logσ2y, gpu_data_handler = get_default_GPUHandler())
-    ζ, logdetΣ = generate_ζ(rng, g, f, ϕ, x, interpreters; n_MC)
-    ζ_cpu = gpu_data_handler(ζ) # differentiable fetch to CPU in Flux package extension
-    #ζi = first(eachcol(ζ_cpu))
-    nLy = reduce(+, map(eachcol(ζ_cpu)) do ζi
-        y_pred_i, logjac = predict_y(ζi, f, transPMs)
-        nLy1 = neg_logden_indep_normal(y_ob, y_pred_i, logσ2y)
-        nLy1 - logjac
-    end) / n_MC
-    ent = entropy_MvNormal(size(ζ, 1), logdetΣ)  # defined in logden_normal
+    ent = entropy_MvNormal(size(ζs, 1), logdetΣ)  # defined in logden_normal
     nLy - ent
 end
 
+"""
+    predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, interpreters;
+        get_transPMs, get_ca_int_PMs, n_sample_pred=200, 
+        gpu_data_handler=get_default_GPUHandler())
 
+Prediction function for hybrid model. Retuns an Array `(n_obs, n_site, n_sample_pred)`.
+"""
+function predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, interpreters;
+    get_transPMs, get_ca_int_PMs, n_sample_pred=200, 
+    gpu_data_handler=get_default_GPUHandler())
+    n_site = size(xM, 2)
+    intm_PMs_gen = get_ca_int_PMs(n_site)
+    tans_PMs_gen = get_transPMs(n_site)
+    ζs, _ = generate_ζ(rng, g, f, CA.getdata(ϕ), CA.getdata(xM),
+    (; interpreters..., PMs = intm_PMs_gen); n_MC = n_sample_pred)
+    ζs_cpu = gpu_data_handler(ζs) #
+    y_pred = stack(map(ζ -> first(predict_y(ζ, f, tans_PMs_gen)), eachcol(ζs_cpu)));
+    y_pred
+end
 
 """
 Generate samples of (inv-transformed) model parameters, ζ, and Log-Determinant
@@ -144,7 +147,7 @@ function _create_random(rng, ::GPUArraysCore.AbstractGPUVector{T}, dims...) wher
     # ignores rng
     # https://discourse.julialang.org/t/help-using-cuda-zygote-and-random-numbers/123458/4?u=bgctw
     # Zygote.@ignore CUDA.randn(rng, dims...)
-    Zygote.@ignore CUDA.randn(dims...)
+    ChainRulesCore.@ignore_derivatives CUDA.randn(dims...)
 end
 
 """ 

src/init_hybrid_params.jl

@@ -0,0 +1,69 @@
+"""
+    init_hybrid_params(θP, θM, ϕg, n_batch; transP=asℝ, transM=asℝ)
+
+Setup ComponentVector of parameters to optimize, and associated tools.
+Returns a NamedTuple of
+- ϕ: A ComponentVector of parameters to optimize
+- transPMs_batch, interpreters: Transformations and interpreters as 
+  required by `neg_elbo_transnorm_gf`.
+- get_transPMs: a function returning transformations `(n_site) -> (;P,Ms)`
+- get_ca_int_PMs: a function returning ComponentArrayInterpreter for PMs vector 
+  with PMs shaped as a matrix of `n_site` columns of `θM`
+
+# Arguments
+- `θP`, `θM`: Template ComponentVectors of global parameters and ML-predicted parameters
+- `ϕg`: vector of parameters to optimize, as returned by `gen_hybridcase_MLapplicator`
+- `n_batch`: the number of sites to predicted in each mini-batch
+- `transP`, `transM`: the Transformations for the global and site-dependent parameters
+"""
+function init_hybrid_params(θP, θM, ϕg, n_batch; transP=asℝ, transM=asℝ)
+    n_θP = length(θP)
+    n_θM = length(θM)
+    n_ϕg = length(ϕg)
+    # zero correlation matrices
+    ρsP = zeros(sum(1:(n_θP - 1)))
+    ρsM = zeros(sum(1:(n_θM - 1)))
+    ϕunc0 = CA.ComponentVector(;
+        logσ2_logP = fill(-10.0, n_θP),
+        coef_logσ2_logMs = reduce(hcat, ([-10.0, 0.0] for _ in 1:n_θM)),
+        ρsP,
+        ρsM)
+    ϕt = CA.ComponentVector(;
+        μP = θP,
+        ϕg = ϕg,
+        unc = ϕunc0);
+    #
+    get_transPMs = let transP=transP, transM=transM, n_θP=n_θP, n_θM=n_θM 
+        function get_transPMs_inner(n_site)
+            transPMs = as(
+                (P = as(Array, transP, n_θP),
+                Ms = as(Array, transM, n_θM, n_site)))
+        end
+    end
+    transPMs_batch = get_transPMs(n_batch)
+    trans_gu = as(
+        (μP = as(Array, asℝ₊, n_θP),
+        ϕg = as(Array, n_ϕg),
+        unc = as(Array, length(ϕunc0))))
+    ϕ = inverse_ca(trans_gu, ϕt)        
+    # trans_g = as(
+    #     (μP = as(Array, asℝ₊, n_θP),
+    #     ϕg = as(Array, n_ϕg)))       
+    #
+    get_ca_int_PMs = let 
+        function get_ca_int_PMs_inner(n_site)
+            ComponentArrayInterpreter(CA.ComponentVector(; θP,
+            θMs = CA.ComponentMatrix(
+                zeros(n_θM, n_site), first(CA.getaxes(θM)), CA.Axis(i = 1:n_site))))
+        end
+        
+    end
+    interpreters = map(get_concrete,
+    (;
+        μP_ϕg_unc = ComponentArrayInterpreter(ϕt),
+        PMs = get_ca_int_PMs(n_batch),
+        unc = ComponentArrayInterpreter(ϕunc0)
+    ))
+    (;ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs)
+end
+