enable independent subranges in parameter s

Author: bgctw

dev/doubleMM.jl

@@ -282,7 +282,7 @@ histogram(vec(mean_y_pred - y_true)) # predictions centered around y_o (or y_tru
 
 # look at θP, θM1 of first site
 intm_PMs_gen = get_ca_int_PMs(n_site)
-ζs, _σ = HVI.generate_ζ(rng, g_flux, f, res.u, xM_gpu,
+ζs, _σ = HVI.generate_ζ(rng, g_flux, res.u, xM_gpu,
     (; interpreters..., PMs = intm_PMs_gen); n_MC = n_sample_pred);
 ζs = ζs |> Flux.cpu;
 θPM = vcat(θP_true, θMs_true[:, 1])

src/HybridProblem.jl

@@ -1,13 +1,14 @@
 struct HybridProblem <: AbstractHybridCase 
     θP
     θM
+    f
+    g
+    ϕg
     transP
     transM
+    cor_starts # = (P=(1,),M=(1,))
     n_covar
     n_batch
-    f
-    g
-    ϕg
     train_loader
     # inner constructor to constrain the types
     function HybridProblem(
@@ -17,15 +18,20 @@ struct HybridProblem <: AbstractHybridCase
         transM::Union{Function, Bijectors.Transform}, 
         transP::Union{Function, Bijectors.Transform}, 
         n_covar::Integer, n_batch::Integer, 
-        train_loader::DataLoader)
-        new(θP, θM, transM, transP, n_covar, n_batch, f, g, ϕg, train_loader)
+        train_loader::DataLoader,
+        cor_starts = (P=(1,), M=(1,)))
+        new(θP, θM, f, g, ϕg, transM, transP, cor_starts, n_covar, n_batch, train_loader)
     end
 end
 
 function get_hybridcase_par_templates(prob::HybridProblem; scenario::NTuple = ())
     (; θP = prob.θP, θM = prob.θM)
 end
 
+function get_hybridcase_transforms(prob::HybridProblem; scenario::NTuple = ())
+    (; transP = prob.transP, transM = prob.transM)
+end
+
 function get_hybridcase_sizes(prob::HybridProblem; scenario::NTuple = ())
     n_θM = length(prob.θM)
     n_θP = length(prob.θP)
@@ -46,6 +52,9 @@ function get_hybridcase_train_dataloader(
     return(prob.train_loader)
 end
 
+function get_hybridcase_cor_starts(prob::HybridProblem; scenario = ())
+    prob.cor_starts
+end
 
 # function get_hybridcase_FloatType(prob::HybridProblem; scenario::NTuple = ()) 
 #     eltype(prob.θM)

src/HybridVariationalInference.jl

@@ -43,7 +43,7 @@ include("util_opt.jl")
 export neg_logden_indep_normal, entropy_MvNormal
 include("logden_normal.jl")
 
-#export - all internal
+export get_ca_starts
 include("cholesky.jl")
 
 export neg_elbo_transnorm_gf, predict_gf

src/cholesky.jl

@@ -252,6 +252,55 @@ function transformU_cholesky1(v::GPUArraysCore.AbstractGPUVector; n=invsumn(leng
     return U
 end
 
+# function transformU_block_cholesky1(v::CA.ComponentVector; 
+#     ns=(invsumn(length(v[k])) + 1 for k in keys(v)) # may pass for efficiency
+#     )
+#     blocks = [transformU_cholesky1(v[k]; n) for (k, n) in zip(keys(v), ns)]
+#     U = _create_blockdiag(v[first(keys(v))], blocks) # v only for dispatch: plain matrix for gpu
+# end
+
+
+"""
+    get_ca_starts(vc::ComponentVector)
+
+Return a tuple with starting positions of components in vc. 
+Useful for providing information on correlactions among subranges in a vector.
+"""
+function get_ca_starts(vc::CA.ComponentVector)
+    (1,  (1 .+ cumsum((length(vc[k]) for k in front(keys(vc)))))...) 
+end
+"omit the last n elements of an iterator"
+front(itr, n=1) = Iterators.take(itr, length(itr)-n)
+
+"""
+    transformU_block_cholesky1(v::AbstractVector, cor_starts = (1,))
+
+Transform a parameterization v of a blockdiagonal of upper triangular matrices
+into the this matrix.
+`cor_starts` is a NTuple of Integeres specifying the first column of each block. 
+E.g. For a matrix with a 3x3, a 2x2, and another block, 
+the blocks start at colums (1,4,6). It defaults to a single entire block.
+"""
+function transformU_block_cholesky1(v::AbstractVector, cor_starts = (1,))
+    cor_starts_end = (cor_starts..., length(v)+1)
+    ranges = ChainRulesCore.@ignore_derivatives (
+        cor_starts_end[i]:(cor_starts_end[i+1]-1) for i in 1:length(cor_starts))
+    blocks = [transformU_cholesky1(v[r]) for r in ranges]
+    U = _create_blockdiag(v, blocks) # v only for dispatch: plain matrix for gpu
+    return(U)
+end
+
+function _create_blockdiag(::AbstractArray, blocks) 
+    BlockDiagonal(blocks)
+end
+
+function _create_blockdiag(::GPUArraysCore.AbstractGPUArray, blocks) 
+    # impose no special structure
+    cat(blocks...; dims=(1, 2))
+end
+
+
+
 () -> begin
     tmp = sqrt.(sum(abs2, U_scaled, dims=1))
     tmp2 = sum(abs2, U_scaled, dims=1) .^ (-1 / 2)

src/elbo.jl

@@ -22,9 +22,9 @@ expected value of the likelihood of observations.
 function neg_elbo_transnorm_gf(rng, g, f, ϕ::AbstractVector, y_ob, xM::AbstractMatrix,
     xP, transPMs, interpreters::NamedTuple; 
     n_MC=3, logσ2y, gpu_data_handler = get_default_GPUHandler(),
-    entropyN = 0.0,
+    cor_starts=(P=(1,),M=(1,))
     )
-    ζs, σ = generate_ζ(rng, g, f, ϕ, xM, interpreters; n_MC)
+    ζs, σ = generate_ζ(rng, g, ϕ, xM, interpreters; n_MC, cor_starts)
     ζs_cpu = gpu_data_handler(ζs) # differentiable fetch to CPU in Flux package extension
     #ζi = first(eachcol(ζs_cpu))
     nLy = reduce(+, map(eachcol(ζs_cpu)) do ζi
@@ -48,13 +48,14 @@ Prediction function for hybrid model. Returns an Array `(n_obs, n_site, n_sample
 """
 function predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, xP, interpreters;
     get_transPMs, get_ca_int_PMs, n_sample_pred=200, 
-    gpu_data_handler=get_default_GPUHandler())
+    gpu_data_handler=get_default_GPUHandler(),
+    cor_starts=(P=(1,),M=(1,)))
     n_site = size(xM, 2)
     intm_PMs_gen = get_ca_int_PMs(n_site)
     trans_PMs_gen = get_transPMs(n_site)
     interpreters_gen = (; interpreters..., PMs = intm_PMs_gen)
-    ζs, _ = generate_ζ(rng, g, f, CA.getdata(ϕ), CA.getdata(xM),
-    interpreters_gen; n_MC = n_sample_pred)
+    ζs, _ = generate_ζ(rng, g, CA.getdata(ϕ), CA.getdata(xM),
+    interpreters_gen; n_MC = n_sample_pred, cor_starts)
     ζs_cpu = gpu_data_handler(ζs) #
     y_pred = stack(map(ζ -> first(predict_y(
         ζ, xP, f, trans_PMs_gen, interpreters_gen.PMs)), eachcol(ζs_cpu)));
@@ -69,14 +70,14 @@ Adds the MV-normally distributed residuals, retrieved by `sample_ζ_norm0`
 to the means extracted from parameters and predicted by the machine learning
 model. 
 """
-function generate_ζ(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix,
-    interpreters::NamedTuple; n_MC=3)
+function generate_ζ(rng, g, ϕ::AbstractVector, xM::AbstractMatrix,
+    interpreters::NamedTuple; n_MC=3, cor_starts=(P=(1,),M=(1,)))
     # see documentation of neg_elbo_transnorm_gf
     ϕc = interpreters.μP_ϕg_unc(CA.getdata(ϕ))
     μ_ζP = ϕc.μP
     ϕg = ϕc.ϕg
     μ_ζMs0 = g(xM, ϕg) # TODO provide μ_ζP to g
-    ζ_resid, σ = sample_ζ_norm0(rng, μ_ζP, μ_ζMs0, ϕc.unc; n_MC)
+    ζ_resid, σ = sample_ζ_norm0(rng, μ_ζP, μ_ζMs0, ϕc.unc; n_MC, cor_starts)
     #ζ_resid, σ = sample_ζ_norm0(rng, ϕ[1:2], reshape(ϕ[2 .+ (1:20)],2,:), ϕ[(end-length(interpreters.unc)+1):end], interpreters.unc; n_MC)
     ζ = stack(map(eachcol(ζ_resid)) do r
         rc = interpreters.PMs(r)
@@ -98,21 +99,21 @@ ComponentMarshellers
 - marsh_batch(n_batch) 
 - marsh_unc(n_UncP, n_UncM, n_UncCorr)
 """
-function sample_ζ_norm0(rng::Random.AbstractRNG, ζP::AbstractVector, ζMs::AbstractMatrix, ϕunc::AbstractVector, args...;
-    n_MC=3) 
+function sample_ζ_norm0(rng::Random.AbstractRNG, ζP::AbstractVector, ζMs::AbstractMatrix, 
+    args...; n_MC, cor_starts)
     n_θP, n_θMs = length(ζP), length(ζMs)
     urand = _create_random(rng, CA.getdata(ζP), n_θP + n_θMs, n_MC)
-    sample_ζ_norm0(urand, ζP, ζMs, ϕunc, args...)
+    sample_ζ_norm0(urand, ζP, ζMs, args...; cor_starts)
 end
 
 function sample_ζ_norm0(urand::AbstractMatrix, ζP::AbstractVector{T}, ζMs::AbstractMatrix, 
-    ϕunc::AbstractVector, int_unc = ComponentArrayInterpreter(ϕunc);
+    ϕunc::AbstractVector, int_unc = ComponentArrayInterpreter(ϕunc); cor_starts
     ) where {T}
     ϕuncc = int_unc(CA.getdata(ϕunc))
     n_θP, n_θMs, (n_θM, n_batch) = length(ζP), length(ζMs), size(ζMs) 
     # make sure to not create a UpperTriangular Matrix of an CuArray in transformU_cholesky1
-    UP = transformU_cholesky1(ϕuncc.ρsP)
-    UM = transformU_cholesky1(ϕuncc.ρsM)
+    UP = transformU_block_cholesky1(ϕuncc.ρsP, cor_starts.P)
+    UM = transformU_block_cholesky1(ϕuncc.ρsM, cor_starts.M)
     cf = ϕuncc.coef_logσ2_logMs
     logσ2_logMs = vec(cf[1, :] .+ cf[2, :] .* ζMs)
     logσ2_logP = vec(CA.getdata(ϕuncc.logσ2_logP))

src/hybrid_case.jl

@@ -11,7 +11,8 @@ For a specific case, provide functions that specify details
 - `get_hybridcase_train_dataloader` (default depends on `gen_hybridcase_synthetic`)
 optionally
 - `gen_hybridcase_synthetic`
-- `get_hybridcase_FloatType` (defaults to eltype(θM))
+- `get_hybridcase_FloatType` (defaults to `eltype(θM)`)
+- `get_hybridcase_cor_starts` (defaults to include all correlations: `(P=(1,), M=(1,))`)
 """
 abstract type AbstractHybridCase end;
 
@@ -93,7 +94,7 @@ function gen_hybridcase_synthetic end
 
 Determine the FloatType for given Case and scenario, defaults to Float32
 """
-function get_hybridcase_FloatType(case::AbstractHybridCase; scenario)
+function get_hybridcase_FloatType(case::AbstractHybridCase; scenario=())
     return eltype(get_hybridcase_par_templates(case; scenario).θM)
 end
 
@@ -114,5 +115,26 @@ function get_hybridcase_train_dataloader(case::AbstractHybridCase, rng::Abstract
     return(train_loader)
 end
 
+"""
+    get_hybridcase_cor_starts(case::AbstractHybridCase; scenario)
+
+Specify blocks in correlation matrices among parameters.
+Returns a NamedTuple.
+- `P`: correlations among global parameters
+- `M`: correlations among ML-predicted parameters
+
+Subsets ofparameters that are correlated with other but not correlated with
+parameters of other subranges are specified by indicating the starting position
+of each subrange.
+E.g. if withing global parameter vector `(p1, p2, p3)`, `p1` and `p2` are correlated, 
+but parameter `p3` is not correlated with them,
+then the first subrange starts at position 1 and the second subrange starts at position 3.
+If there is only single block of all ML-predicted parameters being correlated 
+with each other then this block starts at position 1: `(P=(1,3), M=(1,))`.
+"""
+function get_hybridcase_cor_starts(case::AbstractHybridCase; scenario = ())
+    (P=(1,), M=(1,))
+end
+
 
 

test/test_HybridProblem.jl

@@ -12,13 +12,15 @@ import Zygote
 
 using OptimizationOptimisers
 
+
 const MLengine = Val(nameof(SimpleChains))
 
 construct_problem = () -> begin
     θP = CA.ComponentVector{Float32}(r0 = 0.3, K2 = 2.0)
     θM = CA.ComponentVector{Float32}(r1 = 0.5, K1 = 0.2)
     transP = elementwise(exp)
     transM = Stacked(elementwise(identity), elementwise(exp))
+    cov_starts = (P=(1,2),M=(1)) # assume r0 independent of K2
     n_covar = 5
     n_batch = 10
     int_θdoubleMM = get_concrete(ComponentArrayInterpreter(
@@ -53,7 +55,7 @@ construct_problem = () -> begin
     # HybridProblem(θP, θM, transM, transP, n_covar, n_batch, f_doubleMM_with_global, 
     #     g, ϕg, train_loader)
     HybridProblem(θP, θM, g_chain, f_doubleMM_with_global, 
-        transM, transP, n_covar, n_batch, train_loader)
+        transM, transP, n_covar, n_batch, train_loader, cov_starts)
 end
 prob = construct_problem();
 scenario = (:default,)
@@ -93,3 +95,55 @@ scenario = (:default,)
         @test isapprox(par_templates.θP, int_ϕθP(res.u).θP, rtol = 0.11)
     end
 end
+
+() -> begin
+@testset "neg_elbo_transnorm_gf cpu" begin
+    rng = StableRNG(111)
+    g, ϕg0 = get_hybridcase_MLapplicator(prob, MLengine);
+    train_loader = get_hybridcase_train_dataloader(prob)
+    (xM, xP, y_o) = first(train_loader)
+    n_batch = size(y_o,2)
+    f = get_hybridcase_PBmodel(prob)
+    (θP0, θM0) = get_hybridcase_par_templates(prob)
+    (; transP, transM) = get_hybridcase_transforms(prob)
+
+    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
+        θP0, θM0, ϕg0, n_batch; transP, transM);
+    ϕ_ini = ϕ
+    
+    cost = neg_elbo_transnorm_gf(rng, g, f, ϕ_ini, y_o,
+        xM, xP, transPMs_batch, map(get_concrete, interpreters);
+        n_MC = 8, logσ2y)
+    @test cost isa Float64
+    gr = Zygote.gradient(
+        ϕ -> neg_elbo_transnorm_gf(
+            rng, g, f, ϕ, y_o[:, 1:n_batch],
+            xM[:, 1:n_batch], xP[1:n_batch],
+            transPMs_batch, interpreters; n_MC = 8, logσ2y),
+        CA.getdata(ϕ_ini))
+    @test gr[1] isa Vector
+end;
+
+if CUDA.functional()
+    @testset "neg_elbo_transnorm_gf gpu" begin
+        ϕ = CuArray(CA.getdata(ϕ_ini))
+        xMg_batch = CuArray(xM[:, 1:n_batch])
+        xP_batch = xP[1:n_batch] # used in f which runs on CPU
+        cost = neg_elbo_transnorm_gf(rng, g_flux, f, ϕ, y_o[:, 1:n_batch], 
+            xMg_batch, xP_batch,
+            transPMs_batch, map(get_concrete, interpreters);
+            n_MC = 8, logσ2y)
+        @test cost isa Float64
+        gr = Zygote.gradient(
+            ϕ -> neg_elbo_transnorm_gf(
+                rng, g_flux, f, ϕ, y_o[:, 1:n_batch], 
+                xMg_batch, xP_batch,
+                transPMs_batch, interpreters; n_MC = 8, logσ2y),
+            ϕ)
+        @test gr[1] isa CuVector
+        @test eltype(gr[1]) == FT
+    end
+end
+end #if false
+
+