simplify init_hybrid_paramrs

Author: bgctw

src/HybridSolver.jl

@@ -193,9 +193,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
     n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
     hpints = HybridProblemInterpreters(prob; scenario)
     ϕq = get_hybridproblem_ϕq(prob; scenario)
-    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
-        get_hybridproblem_θP(prob; scenario), pt.θM, cor_ends, ϕg0, hpints; 
-        transP, transM, ϕunc0 = ϕq)
+    (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq)
     int_ϕq = interpreters.ϕq
     int_ϕg_ϕq = interpreters.ϕg_ϕq
     transMs = StackedArray(transM, n_batch)
@@ -233,7 +231,7 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPosteriorS
 
     priors_θP_mean, priors_θMs_mean = construct_priors_θ_mean(
         prob, ϕ0_dev.ϕg, keys(pt.θM), pt.θP, θmean_quant, g_dev, transM, transP;
-        scenario, get_ca_int_PMs, gdevs, pbm_covars)
+        scenario, gdevs, pbm_covars)
 
     loss_elbo = get_loss_elbo(
         g_dev, transP, transMs, f_dev, py;
@@ -349,8 +347,7 @@ function compute_elbo_components(
     g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario)
     ϕq = get_hybridproblem_ϕq(prob; scenario)
     (; transP, transM) = get_hybridproblem_transforms(prob; scenario)
-    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
-        θP, θM, cor_ends, ϕg0, n_batch; transP, transM, ϕunc0 = ϕq)
+    (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq)
     if gdev isa MLDataDevices.AbstractGPUDevice
         ϕ0_dev = gdev(ϕ)
         g_dev = gdev(g) # zygote fails if  gdev is a CPUDevice, although should be non-op
@@ -371,7 +368,7 @@ function compute_elbo_components(
     py = get_hybridproblem_neg_logden_obs(prob; scenario)
     priors_θ_mean = construct_priors_θ_mean(
         prob, ϕ0_dev.ϕg, keys(θM), θP, θmean_quant, g_dev, transM;
-        scenario, get_ca_int_PMs, gdev, cdev, pbm_covars)
+        scenario, gdev, cdev, pbm_covars)
     neg_elbo_gtf_components(
         rng, ϕ0_dev, g_dev, transPMs_batch, f, py, xM, xP, y_o, y_unc, i_sites, interpreters;
         solver.n_MC, solver.n_MC_cap, cor_ends, priors_θ_mean)
@@ -382,7 +379,7 @@ In order to let mean of θ stay close to initial point parameter estimates
 construct a prior on mean θ to a Normal around initial prediction.
 """
 function construct_priors_θ_mean(prob, ϕg, keysθM, θP, θmean_quant, g_dev, transM, transP;
-    scenario::Val{scen}, get_ca_int_PMs, gdevs, pbm_covars,
+    scenario::Val{scen}, gdevs, pbm_covars,
     ) where {scen}
     iszero(θmean_quant) ? ([],[]) :
     begin

src/elbo.jl

@@ -392,9 +392,7 @@ function sample_posterior(rng, prob::AbstractHybridProblem, xM::AbstractMatrix;
     pbm_covar_indices = get_pbm_covar_indices(par_templates.θP, pbm_covars)
     hpints = HybridProblemInterpreters(prob; scenario)
     # TODO remove par_templates from init_hybrid_params - might not get updated θP
-    (; ϕ, transPMs_batch, interpreters, get_transPMs) = init_hybrid_params(
-        get_hybridproblem_θP(prob; scenario), par_templates.θM, cor_ends, ϕg0, hpints; 
-        transP, transM, ϕunc0=ϕq)
+    (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq)
     int_ϕg_ϕq = interpreters.ϕg_ϕq
     int_ϕq = interpreters.ϕq
     transMs = StackedArray(transM, n_batch)        

src/init_hybrid_params.jl

@@ -22,22 +22,14 @@ Returns a NamedTuple of
     because this direction is used much more often.
 - `ϕunc0` initial uncertainty parameters, ComponentVector with format of `init_hybrid_ϕunc.`
 """
-function init_hybrid_params(θP::AbstractVector{FT}, θM::AbstractVector{FT},
-        cor_ends::NamedTuple, ϕg::AbstractVector{FT}, hpints::HybridProblemInterpreters;
-        transP = elementwise(identity), transM = elementwise(identity),
-        ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT))) where {FT}
-    ϕq = ϕunc0
-    n_ϕg = length(ϕg)
+function init_hybrid_params(ϕg::AbstractVector{FT}, ϕq::AbstractVector{FT}) where {FT}
     ϕ = CA.ComponentVector(; ϕg, ϕq)
     interpreters = map(get_concrete,
         (;
             ϕg_ϕq = ComponentArrayInterpreter(ϕ),
-            PMs = get_int_PMst_batch(hpints),
             ϕq = ComponentArrayInterpreter(ϕq)
         ))
-    get_transPMs = transPMs_batch = Val(Symbol("deprecated , use stack_ca_int(intPMs)"))
-    get_ca_int_PMs = Val(Symbol("deprecated , use get_int_PMst_site(HybridProblemInterpreters(prob; scenario))"))
-    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs)
+    (; ϕ, interpreters)
 end
 
 function init_hybrid_params_old(θP::AbstractVector{FT}, θM::AbstractVector{FT},

test/test_elbo.jl

@@ -74,19 +74,14 @@ test_scenario = (scenario) -> begin
     ϕq0 = init_hybrid_ϕq(par_templates.θP, par_templates.θM, transP, cor_ends)
     # ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT))
     # ϕq0 = CP.update_μP_by_θP(ϕunc0, θP_true, transP)
-    hpints = HybridProblemInterpreters(probc; scenario)
-    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
-        θP_true, θMs_true[:, 1], cor_ends, ϕg0, hpints; transP, transM, ϕunc0 = ϕq0)
+    (; ϕ, interpreters) = init_hybrid_params(ϕg0, ϕq0)
     int_ϕq = interpreters.ϕq
     int_ϕg_ϕq = interpreters.ϕg_ϕq
-
-    # @descend_code_warntype init_hybrid_params(θP_true, θMs_true[:, 1], cor_ends, ϕg0, n_batch; transP, transM)
-    # @descend_code_warntype CA.ComponentVector(nt)
     ϕ_ini = ϕ
-    transform_tools = nothing # TODO remove
-    # transform_tools = @inferred CP.setup_transform_ζ(
-    #     transP, transM, get_int_PMst_batch(hpints))
-    int_PMs = get_int_PMst_batch(hpints)
+    # transform_tools = nothing # TODO remove
+    # # transform_tools = @inferred CP.setup_transform_ζ(
+    # #     transP, transM, get_int_PMst_batch(hpints))
+    # int_PMs = get_int_PMst_batch(hpints)
 
     if ggdev isa MLDataDevices.AbstractGPUDevice
         scenario_flux = Val((CP._val_value(scenario)..., :use_Flux, :use_gpu))