Merge pull request #27 from EarthyScience/dev

add Howtos to docu

Author: bgctw

_typos.toml

@@ -1,3 +1,6 @@
+[files]
+extend-exclude = ["docs/src_stash/"]
+
 [default.extend-words]
 SOM = "SOM"
 negLogLik = "negLogLik"

dev/doubleMM.jl

@@ -572,7 +572,7 @@ f_allsites = get_hybridproblem_PBmodel(prob0; scenario, use_all_sites = true)
 trans_mP=StackedArray(transP, size(ζsP, 2))
 trans_mMs=StackedArray(transM, size(ζsMs, 1) * size(ζsMs, 3))
 θsP, θsMs = transform_ζs(ζsP, ζsMs; trans_mP, trans_mMs)
-y = apply_process_model(θsP, θsMs, f, xP) 
+y = f(θsP, θsMs, f, xP) 
 #(; y, θsP, θsMs) = HVI.apply_f_trans(ζsP, ζsMs, f_allsites, xP; transP, transM);
 (y_hmc, θsP_hmc, θsMs_hmc) = (; y, θsP, θsMs);
 

docs/.gitignore

@@ -0,0 +1 @@
+/.quarto/

docs/_quarto.yml

@@ -0,0 +1,9 @@
+project:
+  title: "HybridVariationInference documentation"
+  render: 
+    - src/tutorials/basic_cpu.qmd
+    - src/tutorials/*.qmd
+    
+
+
+

docs/make.jl

@@ -23,8 +23,9 @@ makedocs(;
             #"Test quarto markdown" => "tutorials/test1.md",
         ],
         "How to" => [
-            #".. model independent parameters" => "tutorials/how_to_guides/blocks_corr_site.md",
-            #".. model site-global corr" => "tutorials/how_to_guides/corr_site_global.md",
+            ".. model independent parameters" => "tutorials/blocks_corr.md",
+            ".. model site-global corr" => "tutorials/corr_site_global.md",
+            ".. use GPU" => "tutorials/lux_gpu.md",
         ],
         "Explanation" => [
             #"Theory" => "explanation/theory_hvi.md", TODO activate when paper is published

docs/src/tutorials/Project.toml

@@ -1,11 +1,13 @@
 [deps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DistributionFits = "45214091-1ed4-4409-9bcf-fdb48a05e921"
 HybridVariationalInference = "a108c475-a4e2-4021-9a84-cfa7df242f64"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 PairPlots = "43a3c2be-4208-490b-832a-a21dcd55d7da"
@@ -14,3 +16,4 @@ SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

docs/src/tutorials/basic_cpu.md

@@ -327,7 +327,7 @@ epochs of the optimization.
 ``` julia
 (; probo) = solve(probo_sites, solver; rng,
     callback = callback_loss(100), # output during fitting
-    epochs = 10,
+    epochs = 20,
     #is_inferred = Val(true), # activate type-checks 
 );
 ```

docs/src/tutorials/basic_cpu.qmd

@@ -7,7 +7,7 @@ execute:
   daemon: 3600  
 format:
   commonmark:
-    variant: -raw_html
+    variant: -raw_html+tex_math_dollars
     wrap: preserve    
 bibliography: twutz_txt.bib
 ---
@@ -316,7 +316,25 @@ For the parameters, one row corresponds to
 one site. For the drivers and predictions, one column corresponds to one site.
 
 
-{{< include _pbm_matrix.qmd >}}
+```{julia}
+function f_doubleMM_sites(θc::CA.ComponentMatrix, xPc::CA.ComponentMatrix)
+    # extract several covariates from xP
+    ST = typeof(CA.getdata(xPc)[1:1,:])  # workaround for non-type-stable Symbol-indexing
+    S1 = (CA.getdata(xPc[:S1,:])::ST)   
+    S2 = (CA.getdata(xPc[:S2,:])::ST)
+    #
+    # extract the parameters as row-repeated vectors
+    n_obs = size(S1, 1)
+    VT = typeof(CA.getdata(θc)[:,1])   # workaround for non-type-stable Symbol-indexing
+    (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
+        p1 = CA.getdata(θc[:, par]) ::VT
+        repeat(p1', n_obs)  # matrix: same for each concentration row in S1
+    end
+    #
+    # each variable is a matrix (n_obs x n_site)
+    r0 .+ r1 .* S1 ./ (K1 .+ S1) .* S2 ./ (K2 .+ S2)
+end
+```
 
 Again, the function should not rely on the order of parameters but use symbolic indexing 
 to extract the parameter vectors. For type stability of this symbolic indexing,
@@ -345,7 +363,7 @@ epochs of the optimization.
 ```{julia}
 (; probo) = solve(probo_sites, solver; rng,
     callback = callback_loss(100), # output during fitting
-    epochs = 10,
+    epochs = 20,
     #is_inferred = Val(true), # activate type-checks 
 );
 ```

docs/src/tutorials/blocks_corr.md

@@ -0,0 +1,171 @@
+# How to model indenpendent parameter-blocks in the posterior
+
+
+``` @meta
+CurrentModule = HybridVariationalInference  
+```
+
+This guide shows how to configure independent parameter-blocks in the correlations
+of the posterior.
+
+## Motivation
+
+Modelling all correlations among global and site PBM-parameters respectively
+requires many degrees of freedom.
+
+To decrease the number of parameters to estimate, HVI allows to decompose the
+correlations into independent sub-blocks of parameters.
+
+First load necessary packages.
+
+``` julia
+using HybridVariationalInference
+using ComponentArrays: ComponentArrays as CA
+using Bijectors
+using SimpleChains
+using MLUtils
+using JLD2
+using Random
+using CairoMakie
+using PairPlots   # scatterplot matrices
+```
+
+This tutorial reuses and modifies the fitted object saved at the end of the
+[Basic workflow without GPU](@ref) tutorial.
+
+``` julia
+fname = "intermediate/basic_cpu_results.jld2"
+print(abspath(fname))
+prob = probo_cor = load(fname, "probo");
+```
+
+## Specifying blocks in correlation structure
+
+HVI models the posterior of the parameters at unconstrained scale using a
+multivariate normal distribution. It estimates a parameterization of the
+associated blocks in the correlation matrx and requires a specification
+of the block-structure.
+
+This is done by specifying the positions of the end of the blocks for
+the global (P) and the site-specific parameters (M) respectively using
+a `NamedTuple` of integer vectors.
+
+The defaults specifies a single entry, meaning, there is only one big
+block respectively, spanning all parameters.
+
+``` julia
+cor_ends0 = (P=[length(prob.θP)], M=[length(prob.θM)])
+```
+
+    (P = [1], M = [2])
+
+The following specification models one-entry blocks for each each parameter
+in the correlation block the site parameters, i.e. treating all parameters
+independently with not modelling any correlations between them.
+
+``` julia
+cor_ends = (P=[length(prob.θP)], M=1:length(prob.θM))
+```
+
+    (P = [1], M = 1:2)
+
+## Reinitialize parameters for the posterior approximation.
+
+HVI uses additional fitted parameters to represent the means and the
+covariance matrix of the posterior distribution of model parameters.
+With fewer correlations, also the number of those parameters changes,
+and those parameters must be reinitialized after changing the block structure in
+the correlation matrix.
+
+Here, we obtain construct initial estimates. using [`init_hybrid_ϕunc`](@ref)
+
+``` julia
+ϕunc = init_hybrid_ϕunc(cor_ends, zero(eltype(prob.θM)))
+```
+
+In this two-site parameter case, the the blocked structure saves only one degree of freedom:
+
+``` julia
+length(ϕunc), length(probo_cor.ϕunc)
+```
+
+    (5, 6)
+
+## Update the problem and redo the inversion
+
+``` julia
+prob_ind = HybridProblem(prob; cor_ends, ϕunc)
+```
+
+``` julia
+using OptimizationOptimisers
+import Zygote
+
+solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3)
+
+(; probo) = solve(prob_ind, solver; 
+    callback = callback_loss(100), # output during fitting
+    epochs = 20,
+); probo_ind = probo;
+```
+
+## Compare the correated vs. uncorrelated posterior
+
+First, draw a sample.
+
+``` julia
+n_sample_pred = 400
+(y_cor, θsP_cor, θsMs_cor) = (; y, θsP, θsMs) = predict_hvi(
+  Random.default_rng(), probo_cor; n_sample_pred)
+(y_ind, θsP_ind, θsMs_ind) = (; y, θsP, θsMs) = predict_hvi(
+  Random.default_rng(), probo_ind; n_sample_pred)
+```
+
+``` julia
+i_site = 1
+θ1 = vcat(θsP_ind, θsMs_ind[i_site,:,:])
+θ1_nt = NamedTuple(k => CA.getdata(θ1[k,:]) for k in keys(θ1[:,1])) # 
+plt = pairplot(θ1_nt)
+```
+
+![](blocks_corr_files/figure-commonmark/cell-11-output-1.png)
+
+The corner plot of the independent-parameters estimate shows
+no correlations between site parameters, *r*₁ and *K*₁.
+
+``` julia
+i_out = 4
+fig = Figure(); ax = Axis(fig[1,1], xlabel="mean(y)",ylabel="sd(y)")
+ymean_cor = [mean(y_cor[i_out,s,:]) for s in axes(y_cor, 2)]
+ysd_cor = [std(y_cor[i_out,s,:]) for s in axes(y_cor, 2)]
+scatter!(ax, ymean_cor, ysd_cor, label="correlated") 
+ymean_ind = [mean(y_ind[i_out,s,:]) for s in axes(y_ind, 2)]
+ysd_ind = [std(y_ind[i_out,s,:]) for s in axes(y_ind, 2)]
+scatter!(ax, ymean_ind, ysd_ind, label="independent") 
+axislegend(ax, unique=true)
+fig
+```
+
+![](blocks_corr_files/figure-commonmark/cell-12-output-1.png)
+
+``` julia
+plot_sd_vs_mean = (par) -> begin
+  fig = Figure(); ax = Axis(fig[1,1], xlabel="mean($par)",ylabel="sd($par)")
+  θmean_cor = [mean(θsMs_cor[s,par,:]) for s in axes(θsMs_cor, 1)]
+  θsd_cor = [std(θsMs_cor[s,par,:]) for s in axes(θsMs_cor, 1)]
+  scatter!(ax, θmean_cor, θsd_cor, label="correlated") 
+  θmean_ind = [mean(θsMs_ind[s,par,:]) for s in axes(θsMs_ind, 1)]
+  θsd_ind = [std(θsMs_ind[s,par,:]) for s in axes(θsMs_ind, 1)]
+  scatter!(ax, θmean_ind, θsd_ind, label="independent") 
+  axislegend(ax, unique=true)
+  fig
+end
+plot_sd_vs_mean(:K1)
+```
+
+![](blocks_corr_files/figure-commonmark/cell-13-output-1.png)
+
+The inversion that neglects correlations among site parameters results in
+the same magnitude of estimated uncertainty of predictions.
+However, the uncertainty of the model parameters is severely underestimated
+in this example.