remove MLEngine from get_hybridcase_MLapplicator

better infer it from scenario, although this looses some type-stability.

Author: bgctw

dev/doubleMM.jl

@@ -7,7 +7,7 @@ using Statistics
 using ComponentArrays: ComponentArrays as CA
 
 using SimpleChains
-import Flux # to allow for FluxMLEngine and cpu()
+import Flux 
 using MLUtils
 import Zygote
 
@@ -17,20 +17,22 @@ using Bijectors
 using UnicodePlots
 
 const case = DoubleMM.DoubleMMCase()
-const MLengine = Val(nameof(SimpleChains))
-const FluxMLengine = Val(nameof(Flux))
 scenario = (:default,)
 rng = StableRNG(111)
 
 par_templates = get_hybridcase_par_templates(case; scenario)
 
-(; n_covar, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
+#n_covar = get_hybridcase_n_covar(case; scenario)
+#, n_batch, n_θM, n_θP) = get_hybridcase_sizes(case; scenario)
 
 (; xM, n_site, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
 ) = gen_hybridcase_synthetic(rng, case; scenario);
 
+n_covar = size(xM,1)
+
+
 #----- fit g to θMs_true
-g, ϕg0 = get_hybridcase_MLapplicator(case, MLengine; scenario);
+g, ϕg0 = get_hybridcase_MLapplicator(case; scenario);
 (; transP, transM) = get_hybridcase_transforms(case; scenario)
 
 function loss_g(ϕg, x, g, transM)
@@ -90,6 +92,8 @@ FT = get_hybridcase_float_type(case; scenario)
     θP_true, θMs_true[:, 1], ϕg_opt1, n_batch; transP, transM);
 ϕ_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)
@@ -162,7 +166,8 @@ mean_σ_o_MC = 0.006042
 
 ϕ = CA.getdata(ϕ_ini) |> Flux.gpu;
 xM_gpu = xM |> Flux.gpu;
-g_flux, _ = get_hybridcase_MLapplicator(case, FluxMLengine; scenario);
+scenario_flux = (scenario..., :use_Flux)
+g_flux, _ = get_hybridcase_MLapplicator(case; scenario = scenario_flux);
 
 # otpimize using LUX
 () -> begin
@@ -200,7 +205,7 @@ gr = Zygote.gradient(fcost,
 gr_c = CA.ComponentArray(gr[1] |> Flux.cpu, CA.getaxes(ϕ_ini)...)
 
 train_loader = MLUtils.DataLoader((xM_gpu, xP, y_o, y_unc), batchsize = n_batch)
-#train_loader = get_hybridcase_train_dataloader(case, rng; scenario = (scenario..., :use_flux))
+#train_loader = get_hybridcase_train_dataloader(case, rng; scenario = (scenario..., :use_Flux))
 
 optf = Optimization.OptimizationFunction(
     (ϕ, data) -> begin

ext/HybridVariationalInferenceFluxExt.jl

@@ -35,11 +35,12 @@ end
 #     HybridProblem(θP, θM, g, ϕg, args...; kwargs...)
 # end
 
-function HVI.get_hybridcase_MLapplicator(rng::AbstractRNG, case::HVI.DoubleMM.DoubleMMCase, ::Val{:Flux};
+function HVI.construct_3layer_MLApplicator(
+        rng::AbstractRNG, case::HVI.AbstractHybridCase, ::Val{:Flux};
         scenario::NTuple = ())
     (;θM) = get_hybridcase_par_templates(case; scenario)
     n_out = length(θM)
-    n_covar = 5
+    n_covar = get_hybridcase_n_covar(case; scenario)
     #(; n_covar, n_θM) = get_hybridcase_sizes(case; scenario)
     float_type = get_hybridcase_float_type(case; scenario)
     is_using_dropout = :use_dropout ∈ scenario

ext/HybridVariationalInferenceSimpleChainsExt.jl

@@ -19,8 +19,9 @@ end
 
 HVI.apply_model(app::SimpleChainsApplicator, x, ϕ) = app.m(x, ϕ)
 
-function HVI.get_hybridcase_MLapplicator(rng::AbstractRNG, case::HVI.DoubleMM.DoubleMMCase, ::Val{:SimpleChains};
-        scenario::NTuple=())
+function HVI.construct_3layer_MLApplicator(
+    rng::AbstractRNG, case::HVI.AbstractHybridCase, ::Val{:SimpleChains};
+    scenario::NTuple = ())
     n_covar = get_hybridcase_n_covar(case; scenario)
     FloatType = get_hybridcase_float_type(case; scenario)
     (;θM) = get_hybridcase_par_templates(case; scenario)

src/DoubleMM/f_doubleMM.jl

@@ -91,6 +91,13 @@ function HVI.gen_hybridcase_synthetic(rng::AbstractRNG, case::DoubleMMCase;
     )
 end
 
+function HVI.get_hybridcase_MLapplicator(
+    rng::AbstractRNG, case::HVI.DoubleMM.DoubleMMCase; scenario = ())
+    ml_engine = select_ml_engine(; scenario)
+    construct_3layer_MLApplicator(rng, case, ml_engine; scenario)
+end
+
+
 
 
 

src/HybridProblem.jl

@@ -54,7 +54,7 @@ function get_hybridcase_PBmodel(prob::HybridProblem; scenario::NTuple = ())
     prob.f
 end
 
-function get_hybridcase_MLapplicator(prob::HybridProblem, ml_engine; scenario::NTuple = ());
+function get_hybridcase_MLapplicator(prob::HybridProblem; scenario::NTuple = ());
     prob.g, prob.ϕg
 end
 

src/HybridVariationalInference.jl

@@ -17,6 +17,7 @@ export ComponentArrayInterpreter, flatten1, get_concrete
 include("ComponentArrayInterpreter.jl")
 
 export AbstractModelApplicator, construct_ChainsApplicator
+export construct_3layer_MLApplicator, select_ml_engine
 include("ModelApplicator.jl")
 
 export AbstractGPUDataHandler, NullGPUDataHandler, get_default_GPUHandler

src/ModelApplicator.jl

@@ -36,3 +36,31 @@ end
 # function construct_LuxApplicator end
 
 
+"""
+    construct_3layer_MLApplicator(
+        rng::AbstractRNG, case::HVI.AbstractHybridCase, <ml_engine>;
+        scenario::NTuple = ())
+
+`ml_engine` usually is of type `Val{Symbol}`, e.g. Val(:Flux). See `select_ml_engine`.       
+"""
+function construct_3layer_MLApplicator end
+
+"""
+    select_ml_engine(;scenario)
+
+Returns a value type `Val{:Symbol}` to dispatch on the machine learning engine to use.
+- defaults to `Val(:SimpleChains)`
+- `:use_Lux ∈ scenario -> Val(:Lux)`
+- `:use_Flux ∈ scenario -> Val(:Flux)`
+"""
+function select_ml_engine(;scenario)
+    if :use_Lux ∈ scenario
+        return Val(:Lux)
+    elseif :use_Flux ∈ scenario
+        return Val(:Flux)
+    else
+        # default
+        return Val(:SimpleChains)
+    end
+end
+

src/hybrid_case.jl

@@ -19,22 +19,19 @@ abstract type AbstractHybridCase end;
 
 
 """
-    get_hybridcase_MLapplicator([rng::AbstractRNG,] ::AbstractHybridCase, MLEngine; scenario=())
+    get_hybridcase_MLapplicator([rng::AbstractRNG,] ::AbstractHybridCase; 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 get_hybridcase_MLapplicator end    
 
-function get_hybridcase_MLapplicator(case::AbstractHybridCase, MLEngine; scenario=())
-    get_hybridcase_MLapplicator(Random.default_rng(), case, MLEngine; scenario)
+function get_hybridcase_MLapplicator(case::AbstractHybridCase; scenario=())
+    get_hybridcase_MLapplicator(Random.default_rng(), case; scenario)
 end
 
 """
@@ -138,7 +135,7 @@ function get_hybridcase_train_dataloader(rng::AbstractRNG, case::AbstractHybridC
     scenario = ())
     (; xM, xP, y_o, y_unc) = gen_hybridcase_synthetic(rng, case; scenario)
     n_batch = 10
-    xM_gpu = :use_flux ∈ scenario ? CuArray(xM) : xM
+    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)
 end

test/test_HybridProblem.jl

@@ -12,8 +12,6 @@ import Zygote
 
 using OptimizationOptimisers
 
-const MLengine = Val(nameof(SimpleChains))
-
 construct_problem = () -> begin
     FT = Float32
     θP = CA.ComponentVector{FT}(r0=0.3, K2=2.0)
@@ -63,7 +61,7 @@ scenario = (:default,)
 @testset "loss_gf" begin
     #----------- fit g and θP to y_o
     rng = StableRNG(111)
-    g, ϕg0 = get_hybridcase_MLapplicator(prob, MLengine; scenario)
+    g, ϕg0 = get_hybridcase_MLapplicator(prob; scenario)
     train_loader = get_hybridcase_train_dataloader(rng, prob; scenario)
     (xM, xP, y_o, y_unc) = first(train_loader)
     f = get_hybridcase_PBmodel(prob; scenario)
@@ -101,7 +99,7 @@ import Flux
 
 @testset "neg_elbo_transnorm_gf cpu" begin
     rng = StableRNG(111)
-    g, ϕg0 = get_hybridcase_MLapplicator(prob, MLengine)
+    g, ϕg0 = get_hybridcase_MLapplicator(prob)
     train_loader = get_hybridcase_train_dataloader(rng, prob)
     (xM, xP, y_o, y_unc) = first(train_loader)
     n_batch = size(y_o, 2)

test/test_doubleMM.jl

@@ -11,7 +11,6 @@ import Zygote
 
 using OptimizationOptimisers
 
-const MLengine = Val(nameof(SimpleChains))
 const case = DoubleMM.DoubleMMCase()
 scenario = (:default,)
 
@@ -34,7 +33,7 @@ rng = StableRNG(111)
 end
 
 @testset "loss_g" begin
-    g, ϕg0 = get_hybridcase_MLapplicator(rng, case, MLengine; scenario);
+    g, ϕg0 = get_hybridcase_MLapplicator(rng, case; scenario);
     (;transP, transM) = get_hybridcase_transforms(case; scenario)
 
     function loss_g(ϕg, x, g, transM)
@@ -67,7 +66,7 @@ end
 
 @testset "loss_gf" begin
     #----------- fit g and θP to y_o  (without uncertainty, without transforming θP)
-    g, ϕg0 = get_hybridcase_MLapplicator(case, MLengine; scenario);
+    g, ϕg0 = get_hybridcase_MLapplicator(case; scenario);
     (;transP, transM) = get_hybridcase_transforms(case; scenario)
     f = get_hybridcase_PBmodel(case; scenario)