docu strings of Lux

Author: bgctw

.gitignore

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

ext/HybridVariationalInferenceLuxExt.jl

@@ -7,7 +7,10 @@ using Random
 using StatsFuns: logistic
 
 
-
+# AbstractModelApplicator that stores a Lux.StatefulLuxLayer, so that
+# it can be applied with given inputs and parameters
+# The `int_ϕ` ComponentArrayInterpreter, attaches the correct axes to the
+# supplied parameters, that do not need to keep the Axis information
 struct LuxApplicator{MT, IT} <: AbstractModelApplicator 
     stateful_layer::MT
     int_ϕ::IT
@@ -24,7 +27,7 @@ function HVI.construct_ChainsApplicator(rng::AbstractRNG, m::Chain, float_type=F
 end
 
 function HVI.apply_model(app::LuxApplicator, x, ϕ) 
-    ϕc = app.int_ϕ(ϕ)
+    ϕc = app.int_ϕ(CA.getdata(ϕ))
     app.stateful_layer(x, ϕc)
 end
 

src/HybridSolver.jl

@@ -54,6 +54,8 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
     # Zygote.gradient(ϕ0_dev -> loss_gf(ϕ0_dev, data1...)[1], ϕ0_dev)
     optf = Optimization.OptimizationFunction((ϕ, data) -> loss_gf(ϕ, data...)[1],
         Optimization.AutoZygote())
+    # use CA.getdata(ϕ0_dev), i.e. the plain vector to avoid recompiling for specific CA
+    # loss_gf re-attaches the axes
     optprob = OptimizationProblem(optf, CA.getdata(ϕ0_dev), train_loader_dev)
     res = Optimization.solve(optprob, solver.alg; kwargs...)
     ϕ = intϕ(res.u)

src/ModelApplicator.jl

@@ -11,10 +11,10 @@ a tuple with two components:
 - The applicator 
 - a sample parameter vector (type  depends on the used ML-framework)
 
-Implemented are
-- `construct_SimpleChainsApplicator`
-- `construct_FluxApplicator`
-- `construct_LuxApplicator`
+Implemented overloads of `construct_ChainsApplicator` for layers of 
+- `SimpleChains.SimpleChain`
+- `Flux.Chain`
+- `Lux.Chain`
 """
 abstract type AbstractModelApplicator end
 

src/gf.jl

@@ -138,15 +138,18 @@ Create a loss function for given
 - transM: transforamtion of parameters at unconstrained space
 - f(θMs, θP): mechanistic model 
 - intϕ: interpreter attaching axis with components ϕg and ϕP
-- intP: interpreter attaching axis to ζP = ϕP with components used by f
-- kwargs: additional keyword arguments passed to gf, such as gdev or pbm_covars
+- intP: interpreter attaching axis to ζP = ϕP with components used by f,
+  The default, uses `intϕ(ϕ)` as a template
+- kwargs: additional keyword arguments passed to `gf`, such as `gdev` or `pbm_covars`
 
 The loss function `loss_gf(ϕ, xM, xP, y_o, y_unc, i_sites)` takes   
 - parameter vector ϕ
 - xM: matrix of covariate, sites in the batch are in columns
 - xP: iteration of drivers for each site
 - y_o: matrix of observations, sites in columns
 - y_unc: vector of uncertainty information for each observation
+  Currently, hardcoes squared error loss of `(y_pred .- y_o) ./ σ`, 
+  with `σ = exp.(y_unc ./ 2)`.
 - i_sites: index of sites in the batch
 
 and returns a NamedTuple of