tutorail: save with PBM from DoubleMM

Author: bgctw

.gitignore

@@ -14,4 +14,3 @@ docs/src/**/*_files/
 docs/src/**/*.html
 docs/src/**/*.ipynb
 docs/src/**/*Manifest.toml
-docs/src/tutorials/intermediate/*

docs/make.jl

@@ -1,9 +1,11 @@
 using HybridVariationalInference
+import HybridVariationalInference.DoubleMM
 using Documenter
 
 DocMeta.setdocmeta!(HybridVariationalInference, :DocTestSetup, :(using HybridVariationalInference); recursive=true)
 
 makedocs(;
+    #modules=[HybridVariationalInference, HybridVariationalInference.DoubleMM],
     modules=[HybridVariationalInference],
     authors="Thomas Wutzler <[email protected]> and contributors",
     sitename="HybridVariationalInference.jl",

docs/src/reference/reference_public.md

@@ -11,7 +11,7 @@ i.e. the docstrings.
 
 ``` @autodocs
 Modules = [
-    HybridVariationalInference, 
+    HybridVariationalInference, HybridVariationalInference.DoubleMM
 ]
 Private = false
 ```

docs/src/tutorials/basic_cpu.md

@@ -96,7 +96,7 @@ transM = Stacked(HVI.Exp(), HVI.Exp())
 
 Parameter transformations are specified using the `Bijectors` package.
 Because, `Bijectors.elementwise(exp)`, has problems with automatic differentiation (AD)
-on GPU, we use the non-exported [`Exp`]() wrapper inside `Bijectors.Stacked`.
+on GPU, we use the public but non-exported [`Exp`](@ref) wrapper inside `Bijectors.Stacked`.
 
 ### Prior information on parameters at constrained scale
 
@@ -125,8 +125,8 @@ Here, we use synthetic data generated by the package.
 
 ``` julia
 rng = StableRNG(111)
-scenario = Val((:omit_r0, :covarK2, ))
-(; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario)
+(; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(
+    rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,)))
 ```
 
 Lets look at them.
@@ -335,16 +335,25 @@ epochs of the optimization.
 ## Saving the results
 
 Extracting useful information from the optimized HybridProblem is covered
-in the following tutorial. XXLink
-
+in the following [Inspect results of fitted problem](@ref) tutorial.
 In order to use the results from this tutorial in other tutorials,
 the updated `probo` `HybridProblem` and the interpreters are saved to a JLD2 file.
 
+Before the problem is updated to use the redefinition [`DoubleMM.f_doubleMM_sites`](@ref)
+of the PBM in module `DoubleMM` rather than
+module `Main` to allow for easier reloading with JLD2.
+
+``` julia
+f_batch = PBMPopulationApplicator(DoubleMM.f_doubleMM_sites, n_batch; θP, θM, θFix, xPvec=xP[:,1])
+f_allsites = PBMPopulationApplicator(DoubleMM.f_doubleMM_sites, n_site; θP, θM, θFix, xPvec=xP[:,1])
+probo2 = HybridProblem(probo; f_batch, f_allsites)
+```
+
 ``` julia
 using JLD2
 fname = "intermediate/basic_cpu_results.jld2"
 mkpath("intermediate")
-if probo isa AbstractHybridProblem # do not save on failure above
-    jldsave(fname, false, IOStream; probo, interpreters)
+if probo2 isa AbstractHybridProblem # do not save on failure above
+    jldsave(fname, false, IOStream; probo=probo2, interpreters)
 end
 ```

docs/src/tutorials/basic_cpu.qmd

@@ -31,7 +31,7 @@ using DistributionFits
 
 Next, specify many moving parts of the Hybrid variational inference (HVI)
 
-## The process-based model
+## The process-based model 
 The example process based model (PBM) predicts a double-monod constrained rate
 for different substrate concentrations, `S1`, and `S2`.
 
@@ -101,7 +101,7 @@ transM = Stacked(HVI.Exp(), HVI.Exp())
 
 Parameter transformations are specified using the `Bijectors` package.
 Because, `Bijectors.elementwise(exp)`, has problems with automatic differentiation (AD)
-on GPU, we use the non-exported [`Exp`]() wrapper inside `Bijectors.Stacked`.
+on GPU, we use the public but non-exported [`Exp`](@ref) wrapper inside `Bijectors.Stacked`.
 
 ### Prior information on parameters at constrained scale
 
@@ -130,16 +130,16 @@ Here, we use synthetic data generated by the package.
 
 ```{julia}
 rng = StableRNG(111)
-scenario = Val((:omit_r0, :covarK2, ))
-(; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario)
+(; xM, xP, y_o, y_unc) = gen_hybridproblem_synthetic(
+    rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,)))
 ```
 
 ```{julia}
 #| echo: false
 #| eval: false
 () -> begin
     (; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc) =  
-        gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario)
+    gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,)))
 end
 ```
 
@@ -265,7 +265,7 @@ y1 = f_batch(CA.getdata(θP), CA.getdata(θMs), CA.getdata(x_batch))[2]
     #using Cthulhu
     #@descend_code_warntype f_batch(CA.getdata(θP), CA.getdata(θMs), CA.getdata(x_batch))
     prob0 = HVI.DoubleMM.DoubleMMCase()
-    f_batch0 = get_hybridproblem_PBmodel(prob0; scenario, use_all_sites = false)
+    f_batch0 = get_hybridproblem_PBmodel(prob0; use_all_sites = false)
     y1f = f_batch0(θP, θMs, x_batch)[2]
     y1 .- y1f  # equal
 end
@@ -352,17 +352,26 @@ epochs of the optimization.
 
 ## Saving the results
 Extracting useful information from the optimized HybridProblem is covered
-in the following tutorial. XXLink
-
-In order to use the results from this tutorial in other tutorials, 
+in the following [Inspect results of fitted problem](@ref) tutorial.
+In order to use the results from this tutorial in other tutorials,
 the updated `probo` `HybridProblem` and the interpreters are saved to a JLD2 file.
 
+Before the problem is updated to use the redefinition [`DoubleMM.f_doubleMM_sites`](@ref)
+of the PBM in module `DoubleMM` rather than
+module `Main` to allow for easier reloading with JLD2.
+
+```{julia}
+f_batch = PBMPopulationApplicator(DoubleMM.f_doubleMM_sites, n_batch; θP, θM, θFix, xPvec=xP[:,1])
+f_allsites = PBMPopulationApplicator(DoubleMM.f_doubleMM_sites, n_site; θP, θM, θFix, xPvec=xP[:,1])
+probo2 = HybridProblem(probo; f_batch, f_allsites)
+```
+
 ```{julia}
 using JLD2
 fname = "intermediate/basic_cpu_results.jld2"
 mkpath("intermediate")
-if probo isa AbstractHybridProblem # do not save on failure above
-    jldsave(fname, false, IOStream; probo, interpreters)
+if probo2 isa AbstractHybridProblem # do not save on failure above
+    jldsave(fname, false, IOStream; probo=probo2, interpreters)
 end
 ```
 

docs/src/tutorials/inspect_results.md

@@ -18,28 +18,8 @@ using CairoMakie
 using PairPlots   # scatterplot matrices
 ```
 
-After redefinig the process-based model (currently JLD2 cannot save functions),
-the previously optimized Problem can be loaded.
-
-``` 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
-```
+This tutorial uses the fitted object saved in the
+[Basic workflow without GPU](@ref) tutorial.
 
 ``` julia
 fname = "intermediate/basic_cpu_results.jld2"
@@ -77,7 +57,7 @@ They are ComponentArrays with the parameter dimension names that can be used:
 θsMs[1,:r1,:] # sample of r1 of the first site
 ```
 
-### Corner plots
+## Corner plots
 
 The relation between different variables can be well inspected by
 scatterplot matrices, also called corner plots or pair plots.
@@ -92,7 +72,7 @@ i_site = 1
 plt = pairplot(θ1_nt)
 ```
 
-![](inspect_results_files/figure-commonmark/cell-9-output-1.png)
+![](inspect_results_files/figure-commonmark/cell-8-output-1.png)
 
 The plot shows that parameters for the first site, *K*₁ and *r*₁, are correlated,
 but that we did not model correlation with the global parameter, *K*₂.
@@ -101,7 +81,7 @@ Note that this plots shows only the first out of 800 sites.
 HVI estimated a 1602-dimensional posterior distribution including
 covariances among parameters.
 
-### Expected values and marginal variances
+## Expected values and marginal variances
 
 Lets look at how the estimated uncertainty of a site parameter changes with
 its expected value.
@@ -115,7 +95,7 @@ scatter!(ax, θmean, θsd)
 fig
 ```
 
-![](inspect_results_files/figure-commonmark/cell-11-output-1.png)
+![](inspect_results_files/figure-commonmark/cell-10-output-1.png)
 
 We see that *K*₁ across sites ranges from about 0.18 to 0.25, and that
 its estimated uncertainty is about 0.034, slightly decreasing with the
@@ -157,7 +137,7 @@ scatter!(ax, ymean, ysd)
 fig
 ```
 
-![](inspect_results_files/figure-commonmark/cell-14-output-1.png)
+![](inspect_results_files/figure-commonmark/cell-13-output-1.png)
 
 We see that observed values for associated substrate concentrations range about from
 0.51 to 0.59 with an estimated standard deviation around 0.005 that decreases

docs/src/tutorials/inspect_results.qmd

@@ -29,10 +29,8 @@ using CairoMakie
 using PairPlots   # scatterplot matrices
 ```
 
-After redefinig the process-based model (currently JLD2 cannot save functions),
-the previously optimized Problem can be loaded. 
-
-{{< include _pbm_matrix.qmd >}}
+This tutorial uses the fitted object saved at the end of the 
+[Basic workflow without GPU](@ref) tutorial.
 
 ```{julia}
 fname = "intermediate/basic_cpu_results.jld2"
@@ -43,6 +41,8 @@ probo, interpreters = load(fname, "probo", "interpreters");
 ```{julia}
 #| eval: false
 #| echo: false
+# not necessary any more with DoubleMM.f_doubleMM_sites
+# {{< include _pbm_matrix.qmd >}} 
 # outside notebook, need to reset ModelApplicator, due to fθ defined in Notebook module
 #θFix = CA.ComponentVector{eltype(probo.θP)}(r0=0.3)
 θFix = CA.ComponentVector{eltype(probo.θP)}(

docs/src/tutorials/intermediate/basic_cpu_results.jld2

@@ -14,7 +14,7 @@ using MLDataDevices
 import GPUArraysCore # used in conditional breakpoints
 import StableRNGs
 
-export f_doubleMM, xP_S1, xP_S2
+export f_doubleMM, f_doubleMM_sites, xP_S1, xP_S2
 include("f_doubleMM.jl")
 
 

src/DoubleMM/DoubleMM.jl

@@ -21,6 +21,29 @@ int_xP1 = ComponentArrayInterpreter(CA.ComponentVector(S1 = xP_S1, S2 = xP_S2))
 
 const int_θdoubleMM = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM)))
 
+"""
+    f_doubleMM(θc::CA.ComponentVector{ET}, x) where ET
+
+Example process based model (PBM) predicts a double-monod constrained rate
+for different substrate concentration vectors, `x.S1`, and `x.S2` for a single site.
+θc is a ComponentVector with scalar parameters as components: `r0`, `r1`, `K1`, and `K2`
+
+It predicts a rate for each entry in concentrations:
+`y = r0 .+ r1 .* x.S1 ./ (K1 .+ x.S1) .* x.S2 ./ (K2 .+ x.S2)`.
+
+It is defined as 
+```julia
+function f_doubleMM(θc::ComponentVector{ET}, x) where ET
+    # extract parameters not depending on order, i.e whether they are in θP or θM
+    # r0 = θc[:r0]
+    (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
+        getdata(θc[par])::ET
+    end
+    y = r0 .+ r1 .* x.S1 ./ (K1 .+ x.S1) .* x.S2 ./ (K2 .+ x.S2)
+    return (y)
+end
+```
+"""
 function f_doubleMM(θc::CA.ComponentVector{ET}, x) where ET
     # extract parameters not depending on order, i.e whether they are in θP or θM
     GPUArraysCore.allowscalar() do # index to scalar parameter in parameter vector
@@ -40,13 +63,44 @@ function f_doubleMM(θc::CA.ComponentVector{ET}, x) where ET
     end
 end
 
+"""
+    f_doubleMM_sites(θc::CA.ComponentMatrix, xPc::CA.ComponentMatrix)
+
+Example process based model (PBM) that predicts for a batch of sites.
+
+Arguments
+- `θc`: parameters with one row per site and symbolic column index 
+- `xPc`: model drivers with one column per site, and symbolic row index
+
+Returns a matrix `(n_obs x n_site)` of predictions.
+
+```julia
+function f_doubleMM_sites(θc::ComponentMatrix, xPc::ComponentMatrix)
+    # extract several covariates from xP
+    ST = typeof(getdata(xPc)[1:1,:])  # workaround for non-type-stable Symbol-indexing
+    S1 = (getdata(xPc[:S1,:])::ST)   
+    S2 = (getdata(xPc[:S2,:])::ST)
+    #
+    # extract the parameters as vectors that are row-repeated into a matrix
+    n_obs = size(S1, 1)
+    VT = typeof(getdata(θc)[:,1])   # workaround for non-type-stable Symbol-indexing
+    (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
+        p1 = 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
+```
+"""
 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
+    # extract the parameters as vectors that are row-repeated into a matrix
     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