reorganize AbstractHybridCase

Author: bgctw

dev/doubleMM.jl

@@ -24,100 +24,80 @@ using OptimizationOptimisers
 
 using UnicodePlots
 
-const EX = HybridVariationalInference.DoubleMM
 const case = DoubleMM.DoubleMMCase()
 const MLengine = Val(nameof(SimpleChains))
-scenario=(:default,)
-
+scenario = (:default,)
 rng = StableRNG(111)
 
-(; n_covar_pc, n_covar, n_site, n_batch, n_θM, n_θP) = get_case_sizes(case; scenario)
+par_templates = get_hybridcase_par_templates(case; scenario)
+
+(; n_covar, n_site, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
 
-# const int_θP = ComponentArrayInterpreter(EX.θP)
-# const int_θM = ComponentArrayInterpreter(EX.θM)
+# const int_θP = ComponentArrayInterpreter(par_templates.θP)
+# const int_θM = ComponentArrayInterpreter(par_templates.θM)
 # const int_θPMs_flat = ComponentArrayInterpreter(P = n_θP, Ms = n_θM * n_batch)
-# const int_θ = ComponentArrayInterpreter(CA.ComponentVector(;θP=EX.θP,θM=EX.θM))
+# const int_θ = ComponentArrayInterpreter(CA.ComponentVector(;θP=par_templates.θP,θM=par_templates.θM))
 # # moved to f_doubleMM
 # # const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(;θP,θM)))
 # # const S1 = [1.0, 1.0, 1.0, 0.3, 0.1]
 # # const S2 = [1.0, 3.0, 5.0, 5.0, 5.0]
-# θ = CA.getdata(vcat(EX.θP, EX.θM))
-
-# const int_θPMs = ComponentArrayInterpreter(CA.ComponentVector(;EX.θP,
-#     θMs=CA.ComponentMatrix(zeros(n_θM, n_batch), first(CA.getaxes(EX.θM)), CA.Axis(i=1:n_batch))))
+# θ = CA.getdata(vcat(par_templates.θP, par_templates.θM))
 
-# moved to f_doubleMM
-# gen_q(InteractionsCovCor)
-x_o, θMs_true0 = gen_cov_pred(case, rng; scenario)
-# normalize to be distributed around the prescribed true values
-int_θMs_sites = ComponentArrayInterpreter(EX.θM, (n_site,))
-int_θMs_batch = ComponentArrayInterpreter(EX.θM, (n_batch,))
-θMs_true = int_θMs_sites(scale_centered_at(θMs_true0, EX.θM, 0.1));
+# const int_θPMs = ComponentArrayInterpreter(CA.ComponentVector(;par_templates.θP,
+#     θMs=CA.ComponentMatrix(zeros(n_θM, n_batch), first(CA.getaxes(par_templates.θM)), CA.Axis(i=1:n_batch))))
 
-@test isapprox(vec(mean(CA.getdata(θMs_true); dims=2)), CA.getdata(EX.θM), rtol=0.02)
-@test isapprox(vec(std(CA.getdata(θMs_true); dims=2)), CA.getdata(EX.θM) .* 0.1, rtol=0.02)
+(; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o) = gen_hybridcase_synthetic(
+    case, rng; scenario);
 
+@test isapprox(
+    vec(mean(CA.getdata(θMs_true); dims = 2)), CA.getdata(par_templates.θM), rtol = 0.02)
+@test isapprox(vec(std(CA.getdata(θMs_true); dims = 2)),
+    CA.getdata(par_templates.θM) .* 0.1, rtol = 0.02)
 
 #----- fit g to θMs_true
-g, ϕg0 = gen_g(case, MLengine; scenario)
-n_ϕg = length(ϕg0)
+g, ϕg0 = gen_hybridcase_MLapplicator(case, MLengine; scenario);
 
 function loss_g(ϕg, x, g)
     ζMs = g(x, ϕg) # predict the log of the parameters
-    θMs = exp.(ζMs)   
+    θMs = exp.(ζMs)
     loss = sum(abs2, θMs .- θMs_true)
     return loss, θMs
 end
-loss_g(ϕg0, x_o, g)
-Zygote.gradient(x-> loss_g(x, x_o, g)[1], ϕg0);
+loss_g(ϕg0, xM, g)
+Zygote.gradient(x -> loss_g(x, xM, g)[1], ϕg0);
 
-optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg,x_o, g)[1],
+optf = Optimization.OptimizationFunction((ϕg, p) -> loss_g(ϕg, xM, g)[1],
     Optimization.AutoZygote())
 optprob = Optimization.OptimizationProblem(optf, ϕg0);
-res = Optimization.solve(optprob, Adam(0.02), callback=callback_loss(100), maxiters=600);
+res = Optimization.solve(optprob, Adam(0.02), callback = callback_loss(100), maxiters = 600);
 
 ϕg_opt1 = res.u;
-loss_g(ϕg_opt1, x_o, g)
-scatterplot(vec(θMs_true), vec(loss_g(ϕg_opt1, x_o, g)[2]))
-@test cor(vec(θMs_true), vec(loss_g(ϕg_opt1, x_o, g)[2])) > 0.9
-
-#----------- fit g and θP to y_obs
-f = gen_f(case; scenario)
-y_true = f(EX.θP, θMs_true, zip())[2]
-
-σ_o = 0.01
-#σ_o = 0.002
-y_o = y_true .+ reshape(randn(length(y_true)), size(y_true)...) .* σ_o
-scatterplot(vec(y_true), vec(y_o))
-scatterplot(vec(log.(y_true)), vec(log.(y_o)))
-
-# fit g to log(θ_true) ~ x_o
+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
 
-int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(ϕg=1:length(ϕg0), θP=EX.θP))
-p = p0 = vcat(ϕg0, EX.θP .* 0.9);  # slightly disturb θP_true
-# #p = p0 = vcat(ϕg_opt1, θP_true .* 0.9);  # slightly disturb θP_true
-# p0c = int_ϕθP(p0); 
-# #gf(g,f_doubleMM, x_o, pc.ϕg, pc.θP)[1]
+#----------- fit g and θP to y_o
+f = gen_hybridcase_PBmodel(case; scenario)
 
+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 data for the batches as separate vectors wrapped in a tuple
-train_loader = MLUtils.DataLoader((
-    x_o, 
-    fill((), n_site), # xP
-    y_o
-    ), 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), batchsize = n_batch)
 
-loss_gf = get_loss_gf(g, f, Float32[], int_ϕθP)
+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)
 
-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(EX.θP, int_ϕθP(res.u).θP)
+hcat(par_templates.θP, int_ϕθP(res.u).θP)

ext/HybridVariationalInferenceSimpleChainsExt.jl

@@ -12,10 +12,10 @@ HVI.construct_SimpleChainsApplicator(m::SimpleChain) = SimpleChainsApplicator(m)
 
 HVI.apply_model(app::SimpleChainsApplicator, x, ϕ) = app.m(x, ϕ)
 
-function HVI.gen_g(case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
+function HVI.gen_hybridcase_MLapplicator(case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
         scenario::NTuple=())
-    (;n_covar, n_θM) = get_case_sizes(case; scenario)
-    FloatType = get_case_FloatType(case; scenario)
+    (;n_covar, n_θM) = get_hybridcase_sizes(case; scenario)
+    FloatType = get_hybridcase_FloatType(case; scenario)
     n_out = n_θM
     is_using_dropout = :use_dropout ∈ scenario
     g_chain = if is_using_dropout

src/DoubleMM/f_doubleMM.jl

@@ -16,33 +16,58 @@ function f_doubleMM(θ::AbstractVector)
     return (y)
 end
 
-function HybridVariationalInference.gen_f(::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)
-        pred_global = eltype(pred_sites)[]
-        return pred_global, pred_sites
-    end
+function HybridVariationalInference.get_hybridcase_par_templates(::DoubleMMCase; scenario::NTuple = ())
+    (; θP, θM)
 end
 
-function HybridVariationalInference.get_case_sizes(::DoubleMMCase; scenario = ())
+function HybridVariationalInference.get_hybridcase_sizes(::DoubleMMCase; scenario = ())
     n_covar_pc = 2
     n_covar = n_covar_pc + 3 # linear dependent
     n_site = 10^n_covar_pc
     n_batch = 10
     n_θM = length(θM)
     n_θP = length(θP)
-    (; n_covar_pc, n_covar, n_site, n_batch, n_θM, n_θP)
+    (; n_covar, n_site, n_batch, n_θM, n_θP)
 end
 
-function HybridVariationalInference.get_case_FloatType(::DoubleMMCase; scenario)
+function HybridVariationalInference.gen_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)
+        pred_global = eltype(pred_sites)[]
+        return pred_global, pred_sites
+    end
+end
+
+function HybridVariationalInference.get_hybridcase_FloatType(::DoubleMMCase; scenario)
     return Float32
 end
 
-function HybridVariationalInference.gen_cov_pred(case::DoubleMMCase, rng::AbstractRNG;
+function HybridVariationalInference.gen_hybridcase_synthetic(case::DoubleMMCase, rng::AbstractRNG;
         scenario = ())
-    (; n_covar_pc, n_covar, n_site, n_batch, n_θM, n_θP) = get_case_sizes(case; scenario)
-    FloatType = get_case_FloatType(case; scenario)
-    gen_cov_pred(rng, FloatType, n_covar_pc, n_covar, n_site, n_θM;
+    n_covar_pc = 2
+    (; n_covar, n_site, n_batch, 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)
+    xP = fill((), n_site)
+    y_global_true, y_true = f(θP, θMs_true, zip())
+    σ_o = 0.01
+    #σ_o = 0.002
+    y_global_o = y_global_true .+ randn(rng, size(y_global_true)) .* σ_o
+    y_o = y_true .+ randn(rng, size(y_true)) .* σ_o
+    (;
+        xM,
+        θP_true = θP,
+        θMs_true,
+        xP,
+        y_global_true,
+        y_true,
+        y_global_o,
+        y_o,
+    )
 end

src/HybridVariationalInference.jl

@@ -3,7 +3,7 @@ module HybridVariationalInference
 using ComponentArrays: ComponentArrays as CA
 using Random
 using StatsBase # fit ZScoreTransform
-using Combinatorics # gen_cov_pred/combinations
+using Combinatorics # gen_hybridcase_synthetic/combinations
 
 export ComponentArrayInterpreter, flatten1
 include("ComponentArrayInterpreter.jl")
@@ -12,7 +12,10 @@ export AbstractModelApplicator, construct_SimpleChainsApplicator, construct_Flux
        construct_LuxApplicator
 include("ModelApplicator.jl")
 
-export AbstractHybridCase, gen_g, gen_f, get_case_sizes, get_case_FloatType, gen_cov_pred
+export AbstractHybridCase, gen_hybridcase_MLapplicator, gen_hybridcase_PBmodel, get_hybridcase_sizes, get_hybridcase_FloatType, gen_hybridcase_synthetic,
+       get_hybridcase_par_templates, gen_cov_pred
+include("hybrid_case.jl")
+
 export applyf, gf, get_loss_gf
 include("gf.jl")
 

src/gf.jl

@@ -1,52 +1,3 @@
-"""
-Type to dispatch constructing data and network structures
-for different cases of hybrid problem setups
-"""
-abstract type AbstractHybridCase end;
-
-function get_case_sizes end
-
-"""
-Determine the FloatType for given Case and scenario, defaults to Float32
-"""
-function get_case_FloatType(::AbstractHybridCase; scenario)
-    return Float32
-end
-
-function gen_cov_pred end
-    
-"""
-    gen_g(::AbstractHybridCase, MLEngine, n_covar, n_out; scenario::NTuple=())
-
-Construct the machine learning model fro given problem case and ML-Framework and 
-scenario.
-
-The MLEngine is a value type of a Symbol, usually the name of the module, e.g. 
-`const MLengine = Val(nameof(SimpleChains))`.
-
-returns a Tuple of
-- AbstractModelApplicator
-- initial parameter vector
-"""
-function gen_g end    
-
-"""
-    gen_f(::AbstractHybridCase; scenario::NTuple=())
-
-Construct the process-based model function 
-`f(θP::AbstractVector, θMs::AbstractMatrix, x) -> (AbstractVector, AbstractMatrix)`
-with
-- θP: calibrated parameters that are constant across site
-- θMs: calibrated parameters that vary across sites, with a  column for each site
-- x: drivers, indexed by site
-
-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_f end
-
-
 function applyf(f, θMs::AbstractMatrix, θP::AbstractVector, x)
     # predict several sites with same physical parameters
     yv = map(eachcol(θMs), x) do θM, x_site

src/hybrid_case.jl

@@ -0,0 +1,86 @@
+"""
+Type to dispatch constructing data and network structures
+for different cases of hybrid problem setups
+
+For a specific case, provide functions that specify details
+- get_hybridcase_par_templates
+- get_hybridcase_sizes
+- gen_hybridcase_MLapplicator
+- gen_hybridcase_PBmodel
+optionally
+- gen_hybridcase_synthetic
+- get_hybridcase_FloatType (if it shoudl differ from Float32)
+"""
+abstract type AbstractHybridCase end;
+
+"""
+    get_hybridcase_par_templates(::AbstractHybridCase; scenario)
+
+Provide tuple of templates of ComponentVectors `θP` and `θM`.
+"""
+function get_hybridcase_par_templates end    
+
+"""
+    get_hybridcase_par_templates(::AbstractHybridCase; scenario)
+
+Provide a NamedTuple of number of 
+- n_covar: covariates xM
+- n_site: all sites in the data
+- n_batch: sites in one minibatch during fitting
+- n_θM, n_θP: entries in parameter vectors
+"""
+function get_hybridcase_sizes end
+
+"""
+    gen_hybridcase_MLapplicator(::AbstractHybridCase, MLEngine, n_covar, n_out; scenario=())
+
+Construct the machine learning model fro given problem case and ML-Framework and 
+scenario.
+
+The MLEngine is a value type of a Symbol, usually the name of the module, e.g. 
+`const MLengine = Val(nameof(SimpleChains))`.
+
+returns a Tuple of
+- AbstractModelApplicator
+- initial parameter vector
+"""
+function gen_hybridcase_MLapplicator end    
+
+"""
+    gen_hybridcase_PBmodel(::AbstractHybridCase; scenario::NTuple=())
+
+Construct the process-based model function 
+`f(θP::AbstractVector, θMs::AbstractMatrix, x) -> (AbstractVector, AbstractMatrix)`
+with
+- θP: calibrated parameters that are constant across site
+- θMs: calibrated parameters that vary across sites, with a  column for each site
+- x: drivers, indexed by site
+
+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
+
+"""
+    gen_hybridcase_synthetic(::AbstractHybridCase, rng; scenario)
+
+Setup synthetic data, a NamedTuple of
+- xM: matrix of covariates, with one column per site
+- θP_true: vector global process-model parameters
+- θMs_true: matrix of site-varying process-model parameters, with 
+- xP: Vector of process-model drivers, with an entry per site
+- y_global_true: vector of global observations
+- y_true: matrix of site-specific observations with one column per site
+- y_global_o, y_o: observations with added noise
+"""
+function gen_hybridcase_synthetic end
+
+"""
+    get_hybridcase_FloatType(::AbstractHybridCase; scenario)
+
+Determine the FloatType for given Case and scenario, defaults to Float32
+"""
+function get_hybridcase_FloatType(::AbstractHybridCase; scenario)
+    return Float32
+end