allow specifying initial uncertainty by init_hybrid_ϕunc

Author: bgctw

dev/doubleMM.jl

@@ -55,7 +55,7 @@ scatterplot(θMs_true[2,:], θMs[2,:])
 prob1o.θP
 scatterplot(vec(y_true), vec(y_pred))
 
-# still overestimating θMs
+# still overestimating θMs and θP
 
 () -> begin # with more iterations?
     prob2 = prob1o

src/HybridVariationalInference.jl

@@ -60,7 +60,7 @@ include("cholesky.jl")
 export neg_elbo_transnorm_gf, predict_gf
 include("elbo.jl")
 
-export init_hybrid_params
+export init_hybrid_params, init_hybrid_ϕunc
 include("init_hybrid_params.jl")
 
 export AbstractHybridSolver, HybridPointSolver, HybridPosteriorSolver

src/cholesky.jl

@@ -290,20 +290,19 @@ end
 Return number of correlation coefficients for a correlation matrix of size `(npar x npar)`
 With blocks starting a postions given with tuple `cor_ends`.
 """
-function get_cor_count(cor_ends)
+function get_cor_count(cor_ends::AbstractVector)
     sum(get_cor_counts(cor_ends))
 end
 function get_cor_counts(cor_ends::AbstractVector{T}) where {T}
     isempty(cor_ends) && return (zero(T))
     cnt_blocks = (
         begin
-            cor_start = i == 1 ? one(T) : cor_ends[i-1] + one(T)
-            cor_ends[i] - cor_start
+            i == 1 ? cor_ends[i] : cor_ends[i] - cor_ends[i-1]
         end for i in 1:length(cor_ends)
     )
-    sumn.(cnt_blocks)
+    get_cor_count.(cnt_blocks)
 end
-function get_cor_count(n_par::T) where {T<:Integer}
+function get_cor_count(n_par::T) where T<:Number # <: Integer causes problems with  AD 
     sumn(n_par - one(T))
 end
 
@@ -318,7 +317,7 @@ E.g. For a matrix with a 3x3, a 2x2, and another single-entry block,
 the blocks start at columns (3,5,6). It defaults to a single entire block.
 """
 function transformU_block_cholesky1(
-    v::AbstractVector{T}, cor_ends::AbstractVector{IT}=Int[]) where {T,IT<:Integer}
+    v::AbstractVector{T}, cor_ends::AbstractVector{TI}=Int[]) where {T,TI<:Integer}
     #@show v, cor_ends
     if length(cor_ends) <= 1 # if there is only one block, return it 
         return transformU_cholesky1(v)
@@ -327,7 +326,7 @@ function transformU_block_cholesky1(
     #@show cor_counts
     ranges = ChainRulesCore.@ignore_derivatives (
         begin
-            cor_start = (i == 1 ? 1 : cor_counts[i-1] + one(IT))
+            cor_start = (i == 1 ? one(TI) : cor_counts[i-1] + one(TI))
             cor_start:cor_counts[i]
         end for i in 1:length(cor_counts)
     )

src/init_hybrid_params.jl

@@ -12,41 +12,38 @@ Returns a NamedTuple of
 
 # Arguments
 - `θP`, `θM`: Template ComponentVectors of global parameters and ML-predicted parameters
+- `cor_ends`: NamedTuple with entries, `P`, and `M`, respectively with 
+   integer vectors of ending columns of parameters blocks
 - `ϕg`: vector of parameters to optimize, as returned by `get_hybridproblem_MLapplicator`
 - `n_batch`: the number of sites to predicted in each mini-batch
 - `transP`, `transM`: the Bijector.Transformations for the global and site-dependent 
     parameters, e.g. `Stacked(elementwise(identity), elementwise(exp), elementwise(exp))`.
     Its the transformation froing from unconstrained to constrained space: θ = Tinv(ζ),
     because this direction is used much more often.
+- `ϕunc0` initial uncertainty parameters, ComponentVector wiht format of `init_hybrid_ϕunc.`
 """
-function init_hybrid_params(θP, θM, cor_ends::NamedTuple, ϕg, n_batch; 
-    transP=elementwise(identity), transM=elementwise(identity))
+function init_hybrid_params(θP::AbstractVector{FT}, θM::AbstractVector{FT},
+        cor_ends::NamedTuple, ϕg::AbstractVector{FT}, n_batch;
+        transP = elementwise(identity), transM = elementwise(identity),
+        ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT))) where {FT}
     n_θP = length(θP)
     n_θM = length(θM)
+    @assert cor_ends.P[end] == n_θP
+    @assert cor_ends.M[end] == n_θM
     n_ϕg = length(ϕg)
     # check translating parameters - can match length?
     _ = Bijectors.inverse(transP)(θP)
     _ = Bijectors.inverse(transM)(θM)
-    FT = eltype(θM)
-    # zero correlation matrices
-    # ρsP = zeros(FT, sum(1:(n_θP - 1)))
-    # ρsM = zeros(FT, sum(1:(n_θM - 1)))
-    ρsP = zeros(FT, get_cor_count(cor_ends.P))
-    ρsM = zeros(FT, get_cor_count(cor_ends.M))
-    ϕunc0 = CA.ComponentVector(;
-        logσ2_logP = fill(FT(-10.0), n_θP),
-        coef_logσ2_logMs = reduce(hcat, (FT[-10.0, 0.0] for _ in 1:n_θM)),
-        ρsP,
-        ρsM)
     ϕ = CA.ComponentVector(;
-        μP = apply_preserve_axes(inverse(transP),θP),
+        μP = apply_preserve_axes(inverse(transP), θP),
         ϕg = ϕg,
-        unc = ϕunc0);
+        unc = ϕunc0)
     #
-    get_transPMs = let transP=transP, transM=transM, n_θP=n_θP, n_θM=n_θM 
+    get_transPMs = let transP = transP, transM = transM, n_θP = n_θP, n_θM = n_θM
         function get_transPMs_inner(n_site)
             transMs = ntuple(i -> transM, n_site)
-            ranges = vcat([1:n_θP], [(n_θP + i0*n_θM) .+ (1:n_θM) for i0 in 0:(n_site-1)])
+            ranges = vcat(
+                [1:n_θP], [(n_θP + i0 * n_θM) .+ (1:n_θM) for i0 in 0:(n_site - 1)])
             transPMs = Stacked((transP, transMs...), ranges)
             transPMs
         end
@@ -56,37 +53,54 @@ function init_hybrid_params(θP, θM, cor_ends::NamedTuple, ϕg, n_batch;
     # inv_trans_gu = Stacked(
     #     (inverse(transP), elementwise(identity), elementwise(identity)), values(ranges))
     # ϕ = inv_trans_gu(CA.getdata(ϕt))        
-    get_ca_int_PMs = let 
+    get_ca_int_PMs = let
         function get_ca_int_PMs_inner(n_site)
-            ComponentArrayInterpreter(CA.ComponentVector(; P=θP,
-            Ms = CA.ComponentMatrix(
-                zeros(n_θM, n_site), first(CA.getaxes(θM)), CA.Axis(i = 1:n_site))))
+            ComponentArrayInterpreter(CA.ComponentVector(; P = θP,
+                Ms = CA.ComponentMatrix(
+                    zeros(n_θM, n_site), first(CA.getaxes(θM)), CA.Axis(i = 1:n_site))))
         end
-        
     end
     interpreters = map(get_concrete,
-    (;
-        μP_ϕg_unc = ComponentArrayInterpreter(ϕ),
-        PMs = get_ca_int_PMs(n_batch),
-        unc = ComponentArrayInterpreter(ϕunc0)
-    ))
-    (;ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs)
+        (;
+            μP_ϕg_unc = ComponentArrayInterpreter(ϕ),
+            PMs = get_ca_int_PMs(n_batch),
+            unc = ComponentArrayInterpreter(ϕunc0)
+        ))
+    (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs)
 end
 
-function init_hybrid_ϕunc(logσ2_logP::AbstractVector{FT}, coef_logσ2_logMs, cor_ends;
-    ρ0 = zeros(FT)) where FT
-    
-        n_θP = length(θP)
-        n_θM = length(θM)
-        n_ϕg = length(ϕg)
-        # TODO zero correlation matrices
-        ρsP = zeros(FT, sum(1:(n_θP - 1)))
-        ρsM = zeros(FT, sum(1:(n_θM - 1)))
-        ϕunc0 = CA.ComponentVector(;
-            logσ2_logP = fill(FT(-10.0), n_θP),
-            coef_logσ2_logMs = reduce(hcat, (FT[-10.0, 0.0] for _ in 1:n_θM)),
-            ρsP,
-            ρsM)
-    end
-    
-    
+"""
+    init_hybrid_ϕunc(cor_ends, ρ0=0f0; logσ2_logP, coef_logσ2_logMs, ρsP, ρsM)
+
+Initialize vector of additional parameter of the approximate posterior.
+
+Arguments:
+- `cor_ends`: NamedTuple with entries, `P`, and `M`, respectively with 
+   integer vectors of ending columns of parameters blocks
+- `ρ0`: default entry for ρsP and ρsM, defaults = 0f0.
+- `coef_logσ2_logM`: default column for `coef_logσ2_logMs`, defaults to `[-10.0, 0.0]`
+
+Returns a `ComponentVector` of 
+- `logσ2_logP`: vector of log-variances of ζP (on log scale).
+  defaults to -10
+- `coef_logσ2_logMs`: offset and slope for the log-variances of ζM scaling with 
+   its value given by columns for each parameter in ζM, defaults to `[-10, 0]`
+- `ρsP` and `ρsM`: parameterization of the upper triangular cholesky factor 
+  of the correlation matrices of ζP and ζM, default to all entries `ρ0`, which defaults to zero.
+"""
+function init_hybrid_ϕunc(
+        cor_ends::NamedTuple,
+        ρ0::FT = 0.0f0,
+        coef_logσ2_logM::AbstractVector{FT} = FT[-10.0, 0.0];
+        logσ2_logP::AbstractVector{FT} = fill(FT(-10.0), cor_ends.P[end]),
+        coef_logσ2_logMs::AbstractMatrix{FT} = reduce(
+            hcat, (coef_logσ2_logM for _ in 1:cor_ends.M[end])),
+        ρsP = fill(ρ0, get_cor_count(cor_ends.P)),
+        ρsM = fill(ρ0, get_cor_count(cor_ends.M)),
+) where {FT}
+    CA.ComponentVector(;
+        logσ2_logP,
+        coef_logσ2_logMs,
+        ρsP,
+        ρsM)
+end

src/util_ca.jl

@@ -7,7 +7,7 @@ function cpu_ca end
 # define in FluxExt
 
 function apply_preserve_axes(f, ca::CA.ComponentArray)
-    CA.ComponentArray(f(ca), CA.getaxes(ca))
+    CA.ComponentArray(f(CA.getdata(ca)), CA.getaxes(ca))
 end
 
 

test/test_cholesky_structure.jl

@@ -162,7 +162,7 @@ end;
     U = CP.transformU_block_cholesky1(v, cor_ends)
     @test diag(U' * U) ≈ ones(4)
     @test U[1:3, 4:4] ≈ zeros(3, 1)
-    gr1 = Zygote.gradient(v -> sum(CP.transformU_block_cholesky1(v, cor_ends)), v)[1] # works nice
+    gr1 = Zygote.gradient(v -> sum(CP.transformU_block_cholesky1(v, cor_ends)), v)[1]; # works nice
     # degenerate case of no correlations
     vc0 = CA.ComponentVector{Float32}()
     cor_ends0 = get_ca_ends(vc0)
@@ -171,7 +171,7 @@ end;
     #collect(ns)
     U = CP.transformU_block_cholesky1(CA.getdata(ρ0), cor_ends0)
     @test diag(U) == [1f0]
-    gr1 = Zygote.gradient(v -> sum(CP.transformU_block_cholesky1(ρ0, cor_ends0)), v)[1] # works nice
+    gr1 = Zygote.gradient(v -> sum(CP.transformU_block_cholesky1(ρ0, cor_ends0)), v)[1]; # works nice
 
     if CUDA.functional() # only run the test, if CUDA is working (not on Github ci)
         vc = v_orig = CA.ComponentVector(b1 = CuArray(1.0f0:3.0f0), b2 = CuArray([5.0f0]))

test/test_elbo.jl

@@ -39,6 +39,7 @@ cor_ends = get_hybridproblem_cor_ends(prob; scenario)
 # transP = elementwise(exp)
 # transM = Stacked(elementwise(identity), elementwise(exp))
 #transM = Stacked(elementwise(identity), elementwise(exp), elementwise(exp)) # test mismatch
+ϕunc0 = init_hybrid_ϕunc(cor_ends, zero(FT))
 (; ϕ, transPMs_batch, interpreters, get_transPMs, get_ca_int_PMs) = init_hybrid_params(
     θP_true, θMs_true[:, 1], cor_ends, ϕg0, n_batch; transP, transM);
 ϕ_ini = ϕ