fix error in confusing parameter positions

depending on scenario-templates need to construct different Interpreters passed to forward model

Author: bgctw

.gitignore

@@ -9,3 +9,4 @@ dev/Manifest*.toml
 tmp/
 **/tmp.svg
 dev/intermediate/*
+dev/tmp.pdf

dev/intermediate/probos.jld2

@@ -227,3 +227,11 @@ function get_hybridproblem_cor_ends(prob::AbstractHybridProblem; scenario = ())
     pt = get_hybridproblem_par_templates(prob; scenario)
     (P = [length(pt.θP)], M = [length(pt.θM)])
 end
+
+
+function setup_PBMpar_interpreter(θP, θM, θall = vcat(θP, θM))
+    keys_fixed = ((k for k in keys(θall) if (k ∉ keys(θP)) & (k ∉ keys(θM)))...,)
+    θFix = θall[keys_fixed]
+    intθ = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM, θFix)))
+    intθ, θFix
+end

src/AbstractHybridProblem.jl

@@ -13,10 +13,10 @@ const transMS = Stacked(elementwise(identity), elementwise(exp))
 
 const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM)))
 
-function f_doubleMM(θ::AbstractVector, x)
+function f_doubleMM(θ::AbstractVector, x, intθ)
     # extract parameters not depending on order, i.e whether they are in θP or θM
     y = GPUArraysCore.allowscalar() do 
-        θc = int_θdoubleMM(θ)
+        θc = intθ(θ)
         #using ComponentArrays: ComponentArrays as CA
         #r0, r1, K1, K2 = θc[(:r0, :r1, :K1, :K2)] # does not work on Zygote+GPU
         r0 = θc[:r0]
@@ -30,8 +30,10 @@ end
 
 function HVI.get_hybridproblem_par_templates(::DoubleMMCase; scenario::NTuple = ())
     if (:omit_r0 ∈ scenario)
+        #return ((; θP = θP_nor0, θM, θf = θP[(:K2r)]))
         return ((; θP = θP_nor0, θM))
     end
+    #(; θP, θM, θf = eltype(θP)[])
     (; θP, θM)
 end
 
@@ -74,11 +76,10 @@ function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::NTuple = ()
     )
     #fsite = (θ, x_site) -> f_doubleMM(θ)  # omit x_site drivers
     par_templates = get_hybridproblem_par_templates(prob; scenario)
-    keys_fixed = ((k for k in keys(θall) if
-    (k ∉ keys(par_templates.θP)) & (k ∉ keys(par_templates.θM)))...,)
-    let θFix = gdev(θall[keys_fixed])
+    intθ, θFix = setup_PBMpar_interpreter(par_templates.θP, par_templates.θM, θall)
+    let θFix = gdev(θFix), intθ = get_concrete(intθ)
         function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, x)
-            pred_sites = applyf(f_doubleMM, θMs, θP, θFix, x)
+            pred_sites = applyf(f_doubleMM, θMs, θP, θFix, x, intθ)
             pred_global = eltype(pred_sites)[]
             return pred_global, pred_sites
         end
@@ -101,7 +102,12 @@ const xP_S2 = Float32[1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0]
 # const xP_S2 = Float32[1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0]
 
 HVI.get_hybridproblem_n_covar(prob::DoubleMMCase; scenario) = 5
-HVI.get_hybridproblem_n_site(prob::DoubleMMCase; scenario) = 800
+function HVI.get_hybridproblem_n_site(prob::DoubleMMCase; scenario) 
+    if (:few_sites ∈ scenario)
+         return(100) 
+    end
+    800
+end
 
 function HVI.get_hybridproblem_train_dataloader(prob::DoubleMMCase; scenario = (), 
     n_batch, rng::AbstractRNG = StableRNG(111), kwargs...

src/DoubleMM/f_doubleMM.jl

@@ -42,7 +42,8 @@ export AbstractHybridProblem, get_hybridproblem_MLapplicator, get_hybridproblem_
        get_hybridproblem_priors,
 #update,
        gen_cov_pred,
-       construct_dataloader_from_synthetic
+       construct_dataloader_from_synthetic,
+       setup_PBMpar_interpreter
 include("AbstractHybridProblem.jl")
 
 export HybridProblem

src/HybridVariationalInference.jl

@@ -1,15 +1,15 @@
-function applyf(f, θMs::AbstractMatrix, θP::AbstractVector, θFix::AbstractVector, xP)
+function applyf(f, θMs::AbstractMatrix, θP::AbstractVector, θFix::AbstractVector, xP, args...; kwargs...)
     # predict several sites with same global parameters θP and fixed parameters θFix
     yv = map(eachcol(θMs), xP) do θM, x_site
-        f(vcat(θP, θM, θFix), x_site)
+        f(vcat(θP, θM, θFix), x_site, args...; kwargs...)
     end
     y = stack(yv)
     return(y)
 end
-function applyf(f, θMs::AbstractMatrix, θPs::AbstractMatrix, θFix::AbstractVector, xP)
+function applyf(f, θMs::AbstractMatrix, θPs::AbstractMatrix, θFix::AbstractVector, xP, args...; kwargs...)
     # do not call f with matrix θ, because .* with vectors S1 would go wrong
     yv = map(eachcol(θMs), eachcol(θPs), xP) do θM, θP, xP_site
-        f(vcat(θP, θM, θFix), xP_site)
+        f(vcat(θP, θM, θFix), xP_site, args...; kwargs...)
     end
     y = stack(yv)
     return(y)