provide xP (PBM drivers) as a Matrix with dataloader

to simplify recreating the   tuples of matrices

Author: bgctw

src/AbstractHybridProblem.jl

@@ -179,7 +179,7 @@ function construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHyb
         )
     (; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, prob; scenario)
     n_site = size(xM,2)
-    @assert length(xP) == n_site
+    @assert size(xP,2) == n_site
     @assert size(y_o,2) == n_site
     @assert size(y_unc,2) == n_site
     i_sites = 1:n_site

src/DoubleMM/f_doubleMM.jl

@@ -6,6 +6,11 @@ const θall = vcat(θP, θM)
 
 const θP_nor0 = θP[(:K2,)]
 
+const xP_S1 = Float32[0.5, 0.5, 0.5, 0.5, 0.4, 0.3, 0.2, 0.1]
+const xP_S2 = Float32[1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0]
+
+int_xP1 = ComponentArrayInterpreter(CA.ComponentVector(S1=xP_S1, S2=xP_S2))
+
 # const transP = elementwise(exp)
 # const transM = elementwise(exp)
 
@@ -164,20 +169,29 @@ function HVI.get_hybridproblem_PBmodel(prob::DoubleMMCase; scenario::NTuple = ()
     intθ1, θFix1 = setup_PBMpar_interpreter(par_templates.θP, par_templates.θM, θall)
     θFix = repeat(θFix1', n_site_batch)
     intθ = get_concrete(ComponentArrayInterpreter((n_site_batch,), intθ1))
+    #int_xPb = ComponentArrayInterpreter((n_site_batch,), int_xP1)
     isP = repeat(axes(par_templates.θP,1)', n_site_batch)  
-    let θFix = θFix, θFix_dev = gdev(θFix), intθ = get_concrete(intθ), isP=isP, n_site_batch=n_site_batch
+    let θFix = θFix, θFix_dev = gdev(θFix), intθ = get_concrete(intθ), isP=isP, 
+        n_site_batch=n_site_batch, 
+        #int_xPb=get_concrete(int_xPb),
+        pos_xP = get_positions(int_xP1)
         function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, xP)
-            @assert length(xP) == n_site_batch
+            @assert size(xP,2) == n_site_batch
             @assert size(θMs,2) == n_site_batch
-            # convert vector of tuples to tuple of matricesByRows
-            # need to supply xP as vectorOfTuples to work with DataLoader
-            # k = first(keys(xP[1]))
-            xPM = (; zip(keys(xP[1]), map(keys(xP[1])) do k
-                #stack(map(r -> r[k], xP))' 
-                stack(map(r -> r[k], xP); dims = 1)
-            end)...)
+            # # convert vector of tuples to tuple of matricesByRows
+            # # need to supply xP as vectorOfTuples to work with DataLoader
+            # # k = first(keys(xP[1]))
+            # xPM = (; zip(keys(xP[1]), map(keys(xP[1])) do k
+            #     #stack(map(r -> r[k], xP))' 
+            #     stack(map(r -> r[k], xP); dims = 1)
+            # end)...)
             #xPM = map(transpose, xPM1)
+            #xPc = int_xPb(CA.getdata(xP))
+            #xPM = (S1 = xPc[:,:S1], S2 = xPc[:,:S2]) # problems with Zygote
             # make sure the same order of columns as in intθ
+            # reshape big matrix into NamedTuple of drivers S1 and S2 
+            #   for broadcasting need sites in rows
+            xPM = map(p -> CA.getdata(xP[p,:])', pos_xP)
             θFixd = (θP isa GPUArraysCore.AbstractGPUVector) ? θFix_dev : θFix
             θ = hcat(CA.getdata(θP[isP]), CA.getdata(θMs)', θFixd)
             pred_sites = f_doubleMM(θ, xPM, intθ)'
@@ -202,8 +216,6 @@ end
 #     return Float32
 # end
 
-const xP_S1 = Float32[0.5, 0.5, 0.5, 0.5, 0.4, 0.3, 0.2, 0.1]
-const xP_S2 = Float32[1.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0]
 
 # two observations more?
 # const xP_S1 = Float32[0.5, 0.5, 0.5, 0.5, 0.5, 0.4, 0.3, 0.1]
@@ -242,7 +254,10 @@ function HVI.gen_hybridproblem_synthetic(rng::AbstractRNG, prob::DoubleMMCase;
     # normalize to be distributed around the prescribed true values
     θMs_true = int_θMs_sites(scale_centered_at(θMs_true0, θM, FloatType(0.1)))
     f = get_hybridproblem_PBmodel(prob; scenario, gdev=identity, use_all_sites = true)
-    xP = fill((; S1 = xP_S1, S2 = xP_S2), n_site)
+    #xP = fill((; S1 = xP_S1, S2 = xP_S2), n_site)
+    int_xPn = ComponentArrayInterpreter(int_xP1, (n_site,))
+    xP = int_xPn(vcat(repeat(xP_S1,1,n_site),repeat(xP_S2,1,n_site)))
+    #xP[:S1,:]
     θP = par_templates.θP
     y_global_true, y_true = f(θP, θMs_true, xP)
     σ_o = FloatType(0.01)

src/gf.jl

@@ -1,5 +1,6 @@
 function applyf(f, θMs::AbstractMatrix, θP::AbstractVector, θFix::AbstractVector, xP, args...; kwargs...)
     # predict several sites with same global parameters θP and fixed parameters θFix
+    #θM, x_site = first(zip(eachcol(θMs), xP))
     yv = map(eachcol(θMs), xP) do θM, x_site
         f(vcat(θP, θM, θFix), x_site, args...; kwargs...)
     end

src/logden_normal.jl

@@ -24,7 +24,9 @@ function neg_logden_indep_normal(obs::AbstractArray, μ::AbstractArray, logσ2::
     # optimize argument logσ2 rather than σs for performance
     #nlogL = sum(σfac .* (1/2) .* logσ2 .+ (1/2) .* exp.(.- logσ2) .* abs2.(obs .- μ))
     # specifying logσ2 instead of σ is not transforming a random variable -> no Jacobian
-    nlogL = sum(σfac .* logσ2 .+ abs2.(obs .- μ) .* exp.(.-logσ2)) / 2
+    obs_data = CA.getdata(obs)
+    μ_data = CA.getdata(μ)
+    nlogL = sum(σfac .* logσ2 .+ abs2.(obs_data .- μ_data) .* exp.(.-logσ2)) / convert(eltype(μ),2)
     return (nlogL)
 end
 # function neg_logden_indep_normal(obss::AbstractMatrix, preds::AbstractMatrix, logσ2::AbstractVector; kwargs...)

test/test_HybridProblem.jl

@@ -39,8 +39,10 @@ construct_problem = (;scenario=(:default,)) -> begin
         y = r0 .+ r1 .* x.S1 ./ (K1 .+ x.S1) .* x.S2 ./ (K2 .+ x.S2)
         return (y)
     end
-    function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, x)
-        pred_sites = applyf(f_doubleMM, θMs, θP, CA.ComponentVector{FT}(), x)
+    function f_doubleMM_with_global(θP::AbstractVector, θMs::AbstractMatrix, xP)
+        #Main.@infiltrate_main
+        #first(eachcol(xP))
+        pred_sites = applyf(f_doubleMM, θMs, θP, CA.ComponentVector{FT}(), eachcol(xP))
         pred_global = eltype(pred_sites)[]
         return pred_global, pred_sites
     end
@@ -93,7 +95,7 @@ test_without_flux = (scenario) -> begin
     prob = probc = construct_problem(;scenario);
     #@descend construct_problem(;scenario)
 
-    @testset "n_input and pbm_covars" begin
+    @testset "n_input and pbm_covars  $(last(scenario))" begin
         g, ϕ_g = get_hybridproblem_MLapplicator(prob; scenario);
         if :covarK2 ∈ scenario
             @test g.app.m.inputdim == (static(6),) # 5 + 1 (ncovar + n_pbm)
@@ -104,7 +106,7 @@ test_without_flux = (scenario) -> begin
         end
     end
 
-    @testset "loss_gf" begin
+    @testset "loss_gf  $(last(scenario))" begin
         #----------- fit g and θP to y_o
         rng = StableRNG(111)
         g, ϕg0 = get_hybridproblem_MLapplicator(prob; scenario)
@@ -157,7 +159,7 @@ gdev = gpu_device()
 test_with_flux = (scenario) -> begin
     prob = probc = construct_problem(;scenario);
 
-    @testset "HybridPointSolver" begin
+    @testset "HybridPointSolver $(last(scenario))" begin
         rng = StableRNG(111)
         solver = HybridPointSolver(; alg=Adam(0.02))
         (; ϕ, resopt, probo) = solve(prob, solver; scenario, rng,
@@ -177,7 +179,7 @@ test_with_flux = (scenario) -> begin
         @test ϕ.ϕP.K2 < 1.5 * log(θP.K2)
     end;
 
-    @testset "HybridPosteriorSolver" begin
+    @testset "HybridPosteriorSolver  $(last(scenario))" begin
         rng = StableRNG(111)
         solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3)
         (; ϕ, θP, resopt) = solve(prob, solver; scenario, rng,
@@ -195,7 +197,7 @@ test_with_flux = (scenario) -> begin
     end;
 
     if gdev isa MLDataDevices.AbstractGPUDevice 
-        @testset "HybridPosteriorSolver gpu" begin
+        @testset "HybridPosteriorSolver gpu  $(last(scenario))" begin
             scenf = (scenario..., :use_Flux, :use_gpu, :omit_r0)
             rng = StableRNG(111)
             # here using DoubleMMCase() directly rather than construct_problem
@@ -239,7 +241,7 @@ test_with_flux = (scenario) -> begin
             end
 
         end;
-        @testset "HybridPosteriorSolver also f on gpu" begin
+        @testset "HybridPosteriorSolver also f on gpu  $(last(scenario))" begin
             scenf = (scenario..., :use_Flux, :use_gpu, :omit_r0, :f_on_gpu)
             rng = StableRNG(111)
             probg = HybridProblem(DoubleMM.DoubleMMCase(); scenario = scenf);

test/test_doubleMM.jl

@@ -47,7 +47,7 @@ fneglogden = get_hybridproblem_neg_logden_obs(prob; scenario)
         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)
-    @test size(xP) == (n_site,)
+    @test size(xP) == (16, n_site)
     @test size(y_o) == (8, n_site)
 
     # test same results for same rng
@@ -57,9 +57,10 @@ fneglogden = get_hybridproblem_neg_logden_obs(prob; scenario)
 end
 
 @testset "f_doubleMM_Matrix" begin
-    is = repeat(axes(θP_true, 1)', n_site)
+    is = repeat((1:length(θP_true))', n_site)
     θvec = CA.ComponentVector(P = θP_true, Ms = θMs_true)
-    xPM = map(xP1s -> repeat(xP1s', n_site), xP[1])
+    #xPM = map(xP1s -> repeat(xP1s', n_site), xP[1])
+    xPM = (S1 = CA.getdata(xP[:S1,:])', S2 = CA.getdata(xP[:S2,:])')
     #θ = hcat(θP_true[is], θMs_true')
     intθ1 = get_concrete(ComponentArrayInterpreter(vcat(θP_true, θMs_true[:, 1])))
     #θpos = get_positions(intθ1)
@@ -71,18 +72,18 @@ end
     end
     y = fy(θvec, xPM)
     y_exp = applyf(HVI.DoubleMM.f_doubleMM, θMs_true, θP_true,
-        Vector{eltype(θP_true)}(undef, 0), xP, intθ1)
+        Vector{eltype(θP_true)}(undef, 0), eachcol(xP), intθ1)
     @test y == y_exp'
     ygrad = Zygote.gradient(θv -> sum(fy(θv, xPM)), θvec)[1]
     if gdev isa MLDataDevices.AbstractGPUDevice
         # θg = gdev(θ)
         # xPMg = gdev(xPM)
         # yg = HVI.DoubleMM.f_doubleMM(θg, xPMg, intθ);
-        θvecg = gdev(θvec)
+        θvecg = gdev(θvec); # errors without ";"
         xPMg = gdev(xPM)
         yg = fy(θvecg, xPMg)
         @test cdev(yg) == y_exp'
-        ygradg = Zygote.gradient(θv -> sum(fy(θv, xPMg)), θvecg)[1] # errors without ";"
+        ygradg = Zygote.gradient(θv -> sum(fy(θv, xPMg)), θvecg)[1] 
         @test ygradg isa CA.ComponentArray
         @test CA.getdata(ygradg) isa GPUArraysCore.AbstractGPUArray
         ygradgc = HVI.apply_preserve_axes(cdev, ygradg) # can print the cpu version
@@ -94,7 +95,7 @@ end
 @testset "neg_logden_obs Matrix" begin
     is = repeat(axes(θP_true, 1)', n_site)
     θvec = CA.ComponentVector(P = θP_true, Ms = θMs_true)
-    xPM = map(xP1s -> repeat(xP1s', n_site), xP[1])
+    xPM = (S1 = CA.getdata(xP[:S1,:])', S2 = CA.getdata(xP[:S2,:])')
     #θ = hcat(θP_true[is], θMs_true')
     intθ1 = get_concrete(ComponentArrayInterpreter(vcat(θP_true, θMs_true[:, 1])))
     #θpos = get_positions(intθ1)
@@ -111,13 +112,13 @@ end
         # θg = gdev(θ)
         # xPMg = gdev(xPM)
         # yg = HVI.DoubleMM.f_doubleMM(θg, xPMg, intθ);
-        θvecg = gdev(θvec)
+        θvecg = gdev(θvec);
         xPMg = gdev(xPM)
-        y_og = gdev(y_o)
+        y_og = gdev(y_o);
         y_uncg = gdev(y_unc)
         costg = fcost(θvecg, xPMg, y_og, y_uncg)
         @test costg ≈ cost
-        ygradg = Zygote.gradient(θv -> fcost(θv, xPMg, y_og, y_uncg), θvecg)[1] # errors without ";"
+        ygradg = Zygote.gradient(θv -> fcost(θv, xPMg, y_og, y_uncg), θvecg)[1]; # errors without ";"
         @test ygradg isa CA.ComponentArray
         @test CA.getdata(ygradg) isa GPUArraysCore.AbstractGPUArray
         ygradgc = HVI.apply_preserve_axes(cdev, ygradg) # can print the cpu version

test/test_elbo.jl

@@ -110,13 +110,13 @@ test_scenario = (scenario) -> begin
     @testset "neg_elbo_gtf cpu" begin
         i_sites = 1:n_batch
         cost = neg_elbo_gtf(rng, ϕ_ini, g, transPMs_batch, f, py,
-            xM[:, i_sites], xP[i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites,
+            xM[:, i_sites], xP[:,i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites,
             map(get_concrete, interpreters);
             cor_ends, pbm_covar_indices)
         @test cost isa Float64
         gr = Zygote.gradient(
             ϕ -> neg_elbo_gtf(rng, ϕ, g, transPMs_batch, f, py,
-                xM[:, i_sites], xP[i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites,
+                xM[:, i_sites], xP[:,i_sites], y_o[:, i_sites], y_unc[:, i_sites], i_sites,
                 map(get_concrete, interpreters);
                 cor_ends, pbm_covar_indices),
             CA.getdata(ϕ_ini))
@@ -128,7 +128,7 @@ test_scenario = (scenario) -> begin
             i_sites = 1:n_batch
             ϕ = ggdev(CA.getdata(ϕ_ini))
             xMg_batch = ggdev(xM[:, i_sites])
-            xP_batch = xP[i_sites] # used in f which runs on CPU
+            xP_batch = xP[:,i_sites] # used in f which runs on CPU
             cost = neg_elbo_gtf(rng, ϕ, g_gpu, transPMs_batch, f, py,
                 xMg_batch, xP_batch, y_o[:, i_sites], y_unc[:, i_sites], i_sites,
                 map(get_concrete, interpreters);