Merge pull request #28 from EarthyScience/dev

Account for priors during inversion

Author: bgctw

Project.toml

@@ -12,11 +12,13 @@ CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DistributionFits = "45214091-1ed4-4409-9bcf-fdb48a05e921"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -47,13 +49,15 @@ CommonSolve = "0.2.4"
 ComponentArrays = "0.15.19"
 DistributionFits = "0.3.9"
 Distributions = "0.25.117"
+FillArrays = "1.13.0"
 Flux = "0.14, 0.15, 0.16"
 Functors = "0.4, 0.5"
 GPUArraysCore = "0.1, 0.2"
 LinearAlgebra = "1.10"
 Lux = "1.4.2"
 MLDataDevices = "1.5, 1.6"
 MLUtils = "0.4.5"
+Missings = "1.2.0"
 Optimization = "3.19.3, 4"
 Random = "1.10.0"
 SimpleChains = "0.4"

_typos.toml

@@ -4,3 +4,4 @@ extend-exclude = ["docs/src_stash/"]
 [default.extend-words]
 SOM = "SOM"
 negLogLik = "negLogLik"
+Missings = "Missings"

dev/doubleMM.jl

@@ -32,7 +32,7 @@ cdev = gdev isa MLDataDevices.AbstractGPUDevice ? cpu_device() : identity
 
 #------ setup synthetic data and training data loader
 prob0_ = HybridProblem(DoubleMM.DoubleMMCase(); scenario);
-(; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc
+(; xM, θP_true, θMs_true, xP,  y_true, y_o, y_unc
 ) = gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario);
 n_site, n_batch = get_hybridproblem_n_site_and_batch(prob0_; scenario)
 ζP_true, ζMs_true = log.(θP_true), log.(θMs_true)
@@ -59,7 +59,7 @@ n_epoch = 80
     maxiters = n_batches_in_epoch * n_epoch);
 # update the problem with optimized parameters
 prob0o = prob1o =probo;
-y_pred_global, y_pred, θMs = gf(prob0o; scenario, is_inferred=Val(true));
+ y_pred, θMs = gf(prob0o; scenario, is_inferred=Val(true));
 # @descend_code_warntype gf(prob0o; scenario)
 #@usingany UnicodePlots
 plt = scatterplot(θMs_true'[:, 1], θMs[:, 1]);
@@ -77,7 +77,7 @@ histogram(vec(y_pred) - vec(y_true)) # predictions centered around y_o (or y_tru
     (; ϕ, resopt) = solve(prob0o, solver1; scenario, rng,
         callback = callback_loss(20), maxiters = 400)
     prob1o = HybridProblem(prob0o; ϕg = cpu_ca(ϕ).ϕg, θP = cpu_ca(ϕ).θP)
-    y_pred_global, y_pred, θMs = gf(prob1o, xM, xP; scenario)
+     y_pred, θMs = gf(prob1o, xM, xP; scenario)
     scatterplot(θMs_true[1, :], θMs[1, :])
     scatterplot(θMs_true[2, :], θMs[2, :])
     prob1o.θP
@@ -91,7 +91,7 @@ end
     (; ϕ, resopt) = solve(prob2, solver1; scenario, rng,
         callback = callback_loss(20), maxiters = 600)
     prob2o = HybridProblem(prob2; ϕg = collect(ϕ.ϕg), θP = ϕ.θP)
-    y_pred_global, y_pred, θMs = gf(prob2o, xM, xP)
+     y_pred, θMs = gf(prob2o, xM, xP)
     prob2o.θP
 end
 
@@ -127,7 +127,7 @@ end
     (; ϕ, resopt) = solve(prob3, solver1; scenario, rng,
         callback = callback_loss(50), maxiters = 600)
     prob3o = HybridProblem(prob3; ϕg = cpu_ca(ϕ).ϕg, θP = cpu_ca(ϕ).θP)
-    y_pred_global, y_pred, θMs = gf(prob3o, xM, xP; scenario)
+     y_pred, θMs = gf(prob3o, xM, xP; scenario)
     scatterplot(θMs_true[2, :], θMs[2, :])
     prob3o.θP
     scatterplot(vec(y_true), vec(y_pred))
@@ -173,7 +173,7 @@ solver_post = HybridPosteriorSolver(; alg = OptimizationOptimisers.Adam(0.01), n
     (y1, θsP1, θsMs1) = (y, θsP, θsMs);
 
     () -> begin # prediction with fitted parameters  (should be smaller than mean)
-        y_pred_global, y_pred2, θMs = gf(prob1o, xM, xP; scenario)
+         y_pred2, θMs = gf(prob1o, xM, xP; scenario)
         scatterplot(θMs_true[1, :], θMs[1, :])
         scatterplot(θMs_true[2, :], θMs[2, :])
         hcat(θP_true, θP) # all parameters overestimated
@@ -366,7 +366,7 @@ end
         # ζMs = invt.transM.(θMs_i)
         # _f = get_hybridproblem_PBmodel(probo; scenario)
         # y_site = map(eachcol(θPs), θMs_i) do θP, θM
-        #     y_global, y = _f(θP, reshape(θM, (length(θM), 1)), xP[[i_site]])
+        #       y = _f(θP, reshape(θM, (length(θM), 1)), xP[[i_site]])
         #     y[:,1]
         # end |> stack
         nLs = get_hybridproblem_neg_logden_obs(

docs/make.jl

@@ -23,9 +23,9 @@ makedocs(;
             #"Test quarto markdown" => "tutorials/test1.md",
         ],
         "How to" => [
+            ".. use GPU" => "tutorials/lux_gpu.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

@@ -8,6 +8,7 @@ DistributionFits = "45214091-1ed4-4409-9bcf-fdb48a05e921"
 HybridVariationalInference = "a108c475-a4e2-4021-9a84-cfa7df242f64"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
 PairPlots = "43a3c2be-4208-490b-832a-a21dcd55d7da"

docs/src/tutorials/basic_cpu.md

@@ -104,14 +104,14 @@ HVI is an approximate bayesian analysis and combines prior information on
 the parameters with the model and observed data.
 
 Here, we provide a wide prior by fitting a Lognormal distributions to
-- the mean corresponding to the initial value provided above
-- the 0.95-quantile 3 times the mean
+- the mode corresponding to the initial value provided above
+- the 0.95-quantile 3 times the mode
 using the `DistributionFits.jl` package.
 
 ``` julia
 θall = vcat(θP, θM)
 priors_dict = Dict{Symbol, Distribution}(
-    keys(θall) .=> fit.(LogNormal, θall, QuantilePoint.(θall .* 3, 0.95)))
+    keys(θall) .=> fit.(LogNormal, θall, QuantilePoint.(θall .* 3, 0.95), Val(:mode)))
 ```
 
 ## Observations, model drivers and covariates

docs/src/tutorials/basic_cpu.qmd

@@ -109,14 +109,14 @@ HVI is an approximate bayesian analysis and combines prior information on
 the parameters with the model and observed data.
 
 Here, we provide a wide prior by fitting a Lognormal distributions to 
-- the mean corresponding to the initial value provided above
-- the 0.95-quantile 3 times the mean 
+- the mode corresponding to the initial value provided above
+- the 0.95-quantile 3 times the mode 
 using the `DistributionFits.jl` package.
 
 ```{julia}
 θall = vcat(θP, θM)
 priors_dict = Dict{Symbol, Distribution}(
-    keys(θall) .=> fit.(LogNormal, θall, QuantilePoint.(θall .* 3, 0.95)))
+    keys(θall) .=> fit.(LogNormal, θall, QuantilePoint.(θall .* 3, 0.95), Val(:mode)))
 ```
 
 ## Observations, model drivers and covariates
@@ -138,7 +138,7 @@ rng = StableRNG(111)
 #| echo: false
 #| eval: false
 () -> begin
-    (; xM, θP_true, θMs_true, xP, y_global_true, y_true, y_global_o, y_o, y_unc) =  
+    (; xM, θP_true, θMs_true, xP,  y_true,  y_o, y_unc) =  
     gen_hybridproblem_synthetic(rng, DoubleMM.DoubleMMCase(); scenario=Val((:omit_r0,)))
 end
 ```