implement HybridProblem

Author: bgctw

dev/doubleMM.jl

@@ -12,8 +12,8 @@ using MLUtils
 import Zygote
 
 using CUDA
-using TransformVariables
 using OptimizationOptimisers
+using Bijectors
 using UnicodePlots
 
 const case = DoubleMM.DoubleMMCase()
@@ -24,13 +24,13 @@ rng = StableRNG(111)
 
 par_templates = get_hybridcase_par_templates(case; scenario)
 
-(; n_covar, n_site, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
+(; n_covar, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
 
-(; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, σ_o
+(; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, σ_o
 ) = gen_hybridcase_synthetic(case, rng; scenario);
 
 #----- fit g to θMs_true
-g, ϕg0 = gen_hybridcase_MLapplicator(case, MLengine; scenario);
+g, ϕg0 = get_hybridcase_MLapplicator(case, MLengine; scenario);
 
 function loss_g(ϕg, x, g)
     ζMs = g(x, ϕg) # predict the log of the parameters
@@ -51,7 +51,7 @@ loss_g(ϕg_opt1, xM, g)
 scatterplot(vec(θMs_true), vec(loss_g(ϕg_opt1, xM, g)[2]))
 @test cor(vec(θMs_true), vec(loss_g(ϕg_opt1, xM, g)[2])) > 0.9
 
-f = gen_hybridcase_PBmodel(case; scenario)
+f = get_hybridcase_PBmodel(case; scenario)
 
 #----------- fit g and θP to y_o
 () -> begin
@@ -84,6 +84,9 @@ end
 #---------- HVI
 logσ2y = 2 .* log.(σ_o)
 n_MC = 3
+transP = elementwise(exp)
+transM = Stacked(elementwise(identity), elementwise(exp))
+
 (; ϕ, 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 = ϕ
@@ -188,7 +191,7 @@ end
 
 ϕ = ϕ_ini |> Flux.gpu;
 xM_gpu = xM |> Flux.gpu;
-g_flux, ϕg0_flux_cpu = gen_hybridcase_MLapplicator(case, FluxMLengine; scenario);
+g_flux, ϕg0_flux_cpu = get_hybridcase_MLapplicator(case, FluxMLengine; scenario);
 
 # otpimize using LUX
 () -> begin

ext/HybridVariationalInferenceFluxExt.jl

@@ -25,7 +25,7 @@ function __init__()
     HVI.set_default_GPUHandler(FluxGPUDataHandler())
 end
 
-function HVI.gen_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:Flux};
+function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:Flux};
         scenario::NTuple = ())
     (; n_covar, n_θM) = get_hybridcase_sizes(case; scenario)
     FloatType = get_hybridcase_FloatType(case; scenario)

ext/HybridVariationalInferenceSimpleChainsExt.jl

@@ -12,7 +12,7 @@ HVI.construct_SimpleChainsApplicator(m::SimpleChain) = SimpleChainsApplicator(m)
 
 HVI.apply_model(app::SimpleChainsApplicator, x, ϕ) = app.m(x, ϕ)
 
-function HVI.gen_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
+function HVI.get_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
         scenario::NTuple=())
     (;n_covar, n_θM) = get_hybridcase_sizes(case; scenario)
     FloatType = get_hybridcase_FloatType(case; scenario)

src/DoubleMM/DoubleMM.jl

@@ -1,10 +1,12 @@
 module DoubleMM
 
 using HybridVariationalInference
+using HybridVariationalInference: HybridVariationalInference as HVI
 using ComponentArrays: ComponentArrays as CA
 using Random
 using Combinatorics
 using StatsFuns: logistic
+using Bijectors
 
 
 include("f_doubleMM.jl")

src/DoubleMM/f_doubleMM.jl

@@ -6,6 +6,10 @@ const S2 = [1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0]
 θP = CA.ComponentVector(r0 = 0.3, K2 = 2.0)
 θM = CA.ComponentVector(r1 = 0.5, K1 = 0.2)
 
+transP = elementwise(exp)
+transM = Stacked(elementwise(identity), elementwise(exp))
+
+
 const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM)))
 
 function f_doubleMM(θ::AbstractVector)
@@ -16,21 +20,26 @@ function f_doubleMM(θ::AbstractVector)
     return (y)
 end
 
-function HybridVariationalInference.get_hybridcase_par_templates(::DoubleMMCase; scenario::NTuple = ())
+function HVI.get_hybridcase_par_templates(::DoubleMMCase; scenario::NTuple = ())
     (; θP, θM)
 end
 
-function HybridVariationalInference.get_hybridcase_sizes(::DoubleMMCase; scenario = ())
+function HVI.get_hybridcase_transforms(::AbstractHybridCase; scenario::NTuple = ())
+    (; transP, transM)
+end
+
+function HVI.get_hybridcase_sizes(::DoubleMMCase; scenario = ())
     n_covar_pc = 2
     n_covar = n_covar_pc + 3 # linear dependent
-    n_site = 10^n_covar_pc
+    #n_site = 10^n_covar_pc
     n_batch = 10
     n_θM = length(θM)
     n_θP = length(θP)
-    (; n_covar, n_site, n_batch, n_θM, n_θP)
+    #(; n_covar, n_site, n_batch, n_θM, n_θP)
+    (; n_covar, n_batch, n_θM, n_θP)
 end
 
-function HybridVariationalInference.gen_hybridcase_PBmodel(::DoubleMMCase; scenario::NTuple = ())
+function HVI.get_hybridcase_PBmodel(::DoubleMMCase; scenario::NTuple = ())
     fsite = (θ, x_site) -> f_doubleMM(θ)  # omit x_site drivers
     function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, x)
         pred_sites = applyf(fsite, θMs, θP, x)
@@ -39,21 +48,22 @@ function HybridVariationalInference.gen_hybridcase_PBmodel(::DoubleMMCase; scena
     end
 end
 
-function HybridVariationalInference.get_hybridcase_FloatType(::DoubleMMCase; scenario)
+function HVI.get_hybridcase_FloatType(::DoubleMMCase; scenario)
     return Float32
 end
 
-function HybridVariationalInference.gen_hybridcase_synthetic(case::DoubleMMCase, rng::AbstractRNG;
+function HVI.gen_hybridcase_synthetic(case::DoubleMMCase, rng::AbstractRNG;
         scenario = ())
     n_covar_pc = 2
-    (; n_covar, n_site, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
+    n_site = 200
+    (; n_covar, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
     FloatType = get_hybridcase_FloatType(case; scenario)
     xM, θMs_true0 = gen_cov_pred(rng, FloatType, n_covar_pc, n_covar, n_site, n_θM;
         rhodec = 8, is_using_dropout = false)
     int_θMs_sites = ComponentArrayInterpreter(θM, (n_site,))
     # normalize to be distributed around the prescribed true values
     θMs_true = int_θMs_sites(scale_centered_at(θMs_true0, θM, 0.1))
-    f = gen_hybridcase_PBmodel(case; scenario)
+    f = get_hybridcase_PBmodel(case; scenario)
     xP = fill((), n_site)
     y_global_true, y_true = f(θP, θMs_true, zip())
     σ_o = 0.01
@@ -62,6 +72,7 @@ function HybridVariationalInference.gen_hybridcase_synthetic(case::DoubleMMCase,
     y_o = y_true .+ randn(rng, size(y_true)) .* σ_o
     (;
         xM,
+        n_site,
         θP_true = θP,
         θMs_true,
         xP,

src/HybridProblem.jl

@@ -0,0 +1,45 @@
+struct HybridProblem <: AbstractHybridCase 
+    θP
+    θM
+    transP
+    transM
+    n_covar
+    n_batch
+    f
+    g
+    ϕg
+    # inner constructor to constrain the types
+    function HybridProblem(
+        θP::CA.ComponentVector, θM::CA.ComponentVector, 
+        transM::Union{Function, Bijectors.Transform}, 
+        transP::Union{Function, Bijectors.Transform}, 
+        n_covar::Integer, n_batch::Integer, 
+        f::Function, g::AbstractModelApplicator, ϕg)
+        new(θP, θM, transM, transP, n_covar, n_batch, f, g, ϕg)
+    end
+end
+
+function get_hybridcase_par_templates(prob::HybridProblem; scenario::NTuple = ())
+    (; θP = prob.θP, θM = prob.θM)
+end
+
+function get_hybridcase_sizes(prob::HybridProblem; scenario::NTuple = ())
+    n_θM = length(prob.θM)
+    n_θP = length(prob.θP)
+    (; n_covar=prob.n_covar, n_batch=prob.n_batch, n_θM, n_θP)
+end
+
+function get_hybridcase_PBmodel(prob::HybridProblem; scenario::NTuple = ())
+    prob.f
+end
+
+function get_hybridcase_MLapplicator(prob::HybridProblem, ml_engine; scenario::NTuple = ());
+    prob.g, prob.ϕg
+end
+
+function get_hybridcase_FloatType(prob::HybridProblem; scenario::NTuple = ()) 
+    eltype(prob.θM)
+end
+
+
+

src/HybridVariationalInference.jl

@@ -22,10 +22,13 @@ include("ModelApplicator.jl")
 export AbstractGPUDataHandler, NullGPUDataHandler, get_default_GPUHandler
 include("GPUDataHandler.jl")
 
-export AbstractHybridCase, gen_hybridcase_MLapplicator, gen_hybridcase_PBmodel, get_hybridcase_sizes, get_hybridcase_FloatType, gen_hybridcase_synthetic,
-       get_hybridcase_par_templates, gen_cov_pred
+export AbstractHybridCase, get_hybridcase_MLapplicator, get_hybridcase_PBmodel, get_hybridcase_sizes, get_hybridcase_FloatType, gen_hybridcase_synthetic,
+       get_hybridcase_par_templates, get_hybridcase_transforms, gen_cov_pred
 include("hybrid_case.jl")
 
+export HybridProblem
+include("HybridProblem.jl")
+
 export applyf, gf, get_loss_gf
 include("gf.jl")
 

src/hybrid_case.jl

@@ -4,9 +4,10 @@ for different cases of hybrid problem setups
 
 For a specific case, provide functions that specify details
 - get_hybridcase_par_templates
+- get_hybridcase_transforms
 - get_hybridcase_sizes
-- gen_hybridcase_MLapplicator
-- gen_hybridcase_PBmodel
+- get_hybridcase_MLapplicator
+- get_hybridcase_PBmodel
 optionally
 - gen_hybridcase_synthetic
 - get_hybridcase_FloatType (if it should differ from Float32)
@@ -20,6 +21,16 @@ Provide tuple of templates of ComponentVectors `θP` and `θM`.
 """
 function get_hybridcase_par_templates end    
 
+
+"""
+    get_hybridcase_transforms(::AbstractHybridCase; scenario)
+
+Return a NamedTupe of
+- `transP`: Bijectors.Transform for the global PBM parameters, θP
+- `transM`: Bijectors.Transform for the single-site PBM parameters, θM
+"""
+function get_hybridcase_transforms end
+
 """
     get_hybridcase_par_templates(::AbstractHybridCase; scenario)
 
@@ -32,7 +43,7 @@ Provide a NamedTuple of number of
 function get_hybridcase_sizes end
 
 """
-    gen_hybridcase_MLapplicator(::AbstractHybridCase, MLEngine, n_covar, n_out; scenario=())
+    get_hybridcase_MLapplicator(::AbstractHybridCase, MLEngine, n_covar, n_out; scenario=())
 
 Construct the machine learning model fro given problem case and ML-Framework and 
 scenario.
@@ -44,10 +55,10 @@ returns a Tuple of
 - AbstractModelApplicator
 - initial parameter vector
 """
-function gen_hybridcase_MLapplicator end    
+function get_hybridcase_MLapplicator end    
 
 """
-    gen_hybridcase_PBmodel(::AbstractHybridCase; scenario::NTuple=())
+    get_hybridcase_PBmodel(::AbstractHybridCase; scenario::NTuple=())
 
 Construct the process-based model function 
 `f(θP::AbstractVector, θMs::AbstractMatrix, x) -> (AbstractVector, AbstractMatrix)`
@@ -60,7 +71,7 @@ returns a tuple of predictions with components
 - first, those that are constant across sites
 - second, those that vary across sites, with a column for each site
 """
-function gen_hybridcase_PBmodel end
+function get_hybridcase_PBmodel end
 
 """
     gen_hybridcase_synthetic(::AbstractHybridCase, rng; scenario)
@@ -84,3 +95,5 @@ Determine the FloatType for given Case and scenario, defaults to Float32
 function get_hybridcase_FloatType(::AbstractHybridCase; scenario)
     return Float32
 end
+
+

src/init_hybrid_params.jl

@@ -12,7 +12,7 @@ Returns a NamedTuple of
 
 # Arguments
 - `θP`, `θM`: Template ComponentVectors of global parameters and ML-predicted parameters
-- `ϕg`: vector of parameters to optimize, as returned by `gen_hybridcase_MLapplicator`
+- `ϕg`: vector of parameters to optimize, as returned by `get_hybridcase_MLapplicator`
 - `n_batch`: the number of sites to predicted in each mini-batch
 - `transP`, `transM`: the Bijector.Transformations for the global and site-dependent 
     parameters, e.g. `Stacked(elementwise(identity), elementwise(exp), elementwise(exp))`.

test/runtests.jl

@@ -13,6 +13,8 @@ const GROUP = get(ENV, "GROUP", "All") # defined in in CI.yml
         @time @safetestset "test_logden_normal" include("test_logden_normal.jl")
         #@safetestset "test" include("test/test_doubleMM.jl")
         @time @safetestset "test_doubleMM" include("test_doubleMM.jl")
+        #@safetestset "test" include("test/test_HybridProblem.jl")
+        @time @safetestset "test_HybridProblem" include("test_HybridProblem.jl")
         #@safetestset "test" include("test/test_cholesky_structure.jl")
         @time @safetestset "test_cholesky_structure" include("test_cholesky_structure.jl")
         #@safetestset "test" include("test/test_sample_zeta.jl")