HybridVariationalInference

AbstractComponentArrayInterpreter

Interface for Type that implements

• as_ca(::AbstractArray, interpreter) -> ComponentArray
• Base.length(interpreter) -> Int

When called on a vector, forwards to as_ca.

Type to dispatch constructing data and network structures for different cases of hybrid problem setups.

For a specific prob, provide functions that specify details

• get_hybridproblem_MLapplicator
• get_hybridproblem_transforms
• get_hybridproblem_PBmodel
• get_hybridproblem_neg_logden_obs
• get_hybridproblem_par_templates
• get_hybridproblem_ϕunc
• get_hybridproblem_train_dataloader (default depends on gen_hybridcase_synthetic)
• get_hybridproblem_priors

optionally

• gen_hybridcase_synthetic
• get_hybridproblem_n_covar (defaults to number of rows in xM in train_dataloader )
• get_hybridproblem_float_type (defaults to eltype(θM))
• get_hybridproblem_cor_ends (defaults to include all correlations: (P = [length(θP)], M = [length(θM)]))

The initial value of parameters to estimate is spread

• ϕg: parameter of the MLapplicator: returned by get_hybridproblem_MLapplicator
• ζP: mean of the PBmodel parameters: returned by get_hybridproblem_par_templates
• ϕunc: additional parameters of the approximte posterior: returned by get_hybridproblem_ϕunc

AbstractModelApplicator(x, ϕ)

Abstraction of applying a machine learning model at covariate matrix, x, using parameters, ϕ. It returns a matrix of predictions with the same number of rows as in x.

Constructors for specifics are defined in extension packages. Each constructor takes a special type of machine learning model and returns 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

Non-Concrete version of AbstractComponentArrayInterpreter that avoids storing additional type parameters.

Does not trigger specialization for Interpreters of different axes, but does not allow compiler-inferred length to construct StaticArrays.

Use get_concrete(cai::ComponentArrayInterpreter) to pass a concrete version to performance-critical functions.

ComponentArrayInterpreter(; kwargs...)
ComponentArrayInterpreter(::AbstractComponentArray)
ComponentArrayInterpreter(::AbstractComponentArray, n_dims::NTuple{N,<:Integer})
ComponentArrayInterpreter(n_dims::NTuple{N,<:Integer}, ::AbstractComponentArray)

Construct a ComponentArrayInterpreter <: AbstractComponentArrayInterpreter with components being vectors of given length or given model of a AbstractComponentArray.

The other constructors allow constructing arrays with additional dimensions.

'''julia interpreter = ComponentArrayInterpreter(; P=2, M=(2,3), Unc=5) v = 1.0:length(interpreter) interpreter(v).M == 2 .+ [1 3 5; 2 4 6] vm = stack((v,v .* 10, v .* 100))

intm = ComponentArrayInterpreter(interpreter(v), (3,))
intm(vm)[:Unc, 2]

'''

ComponentArrayInterpreter(; kwargs...)
ComponentArrayInterpreter(::AbstractComponentArray)
ComponentArrayInterpreter(::AbstractComponentArray, n_dims::NTuple{N,<:Integer})
ComponentArrayInterpreter(n_dims::NTuple{N,<:Integer}, ::AbstractComponentArray)

Construct a ComponentArrayInterpreter <: AbstractComponentArrayInterpreter with components being vectors of given length or given model of a AbstractComponentArray.

The other constructors allow constructing arrays with additional dimensions.

'''julia interpreter = ComponentArrayInterpreter(; P=2, M=(2,3), Unc=5) v = 1.0:length(interpreter) interpreter(v).M == 2 .+ [1 3 5; 2 4 6] vm = stack((v,v .* 10, v .* 100))

intm = ComponentArrayInterpreter(interpreter(v), (3,))
intm(vm)[:Unc, 2]

'''

MagnitudeModelApplicator(app, y0)

Wrapper around AbstractModelApplicator that multiplies the prediction by the absolute inverse of an initial estimate of the prediction.

This helps to keep raw predictions and weights in a similar magnitude.

NormalScalingModelApplicator(app, μ, σ)
NormalScalingModelApplicator(app, priors, transM)

Wrapper around AbstractModelApplicator that transforms each output (assumed in [0..1], usch as output of logistic activation function) to a quantile of a Normal distribution.

Length of μ, σ must correspond to the number of outputs of the wrapped ModelApplicator.

This helps to keep raw ML-predictions (in confidence bounds) and weights in a similar magnitude. Compared to specifying bounds, this allows for the possibility (although harder to converge) far beyond the confidence bounds.

The second constructor fits a normal distribution of the inverse-transformed 5% and 95% quantiles of prior distributions.

NullModelApplicator()

Model applicator that returns its inputs. Used for testing.

Concrete version of AbstractComponentArrayInterpreter that stores an axis in its type signature.

Allows compiler-inferred length to construct StaticArrays, but requires specialization on dispatch when provided as an argument to a function.

Use get_concrete(cai::ComponentArrayInterpreter) to pass a concrete version to performance-critical functions.

as_ca(v::AbstractArray, interpretor)

Returns a ComponentArray with underlying data v.

create a function (state, l) -> false that prints iter and loss each moditer

Create n_covar correlated covariates from uncorrelated row-wise vector x_pc, with correlations rhos to the linear combinations of x_pc.

By default correlations, rhos = (1.0),0.88,0.78,0.69,0.61 ..., decrease exponentially as e^{-i/rhodec}, with rhodec = 8.

construct_3layer_MLApplicator(
    rng::AbstractRNG, prob::HVI.AbstractHybridProblem, <ml_engine>;
    scenario::NTuple = ())

ml_engine usually is of type Val{Symbol}, e.g. Val(:Flux). See select_ml_engine.

construct_ChainsApplicator([rng::AbstractRNG,] chain, float_type)

cpu_ca(ca::CA.ComponentArray)

Move ComponentArray form gpu to cpu.

Removes the highest level of keys. Keeps the reference to the underlying data, but changes the axis. If first-level vector has no sub-names, an error (Aguement Error tuple must be non-empty) is thrown.

omit the last n elements of an iterator

Generate correlated covariates and synthetic true parameters that are a linear combination of the uncorrelated underlying principal factors and their binary combinations.

In addition provide a SimpleChains model of adequate complexity to fit this relationship θMstrue = f(xo)

gen_hybridcase_synthetic([rng,] ::AbstractHybridProblem; scenario)

Setup synthetic data, a NamedTuple of

• xM: matrix of covariates, with one column per site
• θP_true: vector global process-model parameters
• θMs_true: matrix of site-varying process-model parameters, with
• xP: Vector of process-model drivers, with an entry per site
• yglobaltrue: vector of global observations
• y_true: matrix of site-specific observations with one column per site
• yglobalo, y_o: observations with added noise

Generate samples of (inv-transformed) model parameters, ζ, and the vector of standard deviations, σ, i.e. the diagonal of the cholesky-factor.

Adds the MV-normally distributed residuals, retrieved by sample_ζ_norm0 to the means extracted from parameters and predicted by the machine learning model.

get_ca_ends(vc::ComponentVector)

Return a Vector with ending positions of components in vc. Useful for providing information on correlactions among subranges in a vector.

get_ca_starts(vc::ComponentVector)

Return a tuple with starting positions of components in vc. Useful for providing information on correlactions among subranges in a vector.

get_cor_count(cor_ends::AbstractVector)
get_cor_count(n_par::Integer)

Return number of correlation coefficients for a correlation matrix of size (npar x npar) With blocks starting a positions given with tuple cor_ends.

get_hybridproblem_MLapplicator([rng::AbstractRNG,] ::AbstractHybridProblem; scenario=())

Construct the machine learning model fro given problem prob and ML-Framework and scenario.

returns a Tuple of

• AbstractModelApplicator
• initial parameter vector

get_hybridproblem_PBmodel(::AbstractHybridProblem; scenario::NTuple=())

Construct the process-based model function f(θP::AbstractVector, θMs::AbstractMatrix, x) -> (AbstractVector, AbstractMatrix) with

• θP: calibrated parameters that are constant across site
• θMs: calibrated parameters that vary across sites, with a column for each site
• x: drivers, indexed by site

returns a tuple of predictions with components

• first, those that are constant across sites
• second, those that vary across sites, with a column for each site

get_hybridproblem_cor_ends(prob::AbstractHybridProblem; scenario)

Specify blocks in correlation matrices among parameters. Returns a NamedTuple.

• P: correlations among global parameters
• M: correlations among ML-predicted parameters

Subsets ofparameters that are correlated with other but not correlated with parameters of other subranges are specified by indicating the starting position of each subrange. E.g. if within global parameter vector (p1, p2, p3), p1 and p2 are correlated, but parameter p3 is not correlated with them, then the first subrange starts at position 1 and the second subrange starts at position 3. If there is only single block of all ML-predicted parameters being correlated with each other then this block starts at position 1: (P=(1,3), M=(1,)).

get_hybridproblem_float_type(::AbstractHybridProblem; scenario)

Determine the FloatType for given Case and scenario, defaults to Float32

get_hybridproblem_n_covar(::AbstractHybridProblem; scenario)

Provide the number of covariates. Default returns the number of rows in xM from get_hybridproblem_train_dataloader.

get_hybridproblem_neg_logden_obs(::AbstractHybridProblem; scenario)

Provide a function(y_obs, ypred) -> Real that computes the negative logdensity of the observations, given the predictions.

get_hybridproblem_par_templates(::AbstractHybridProblem; scenario)

Provide tuple of templates of ComponentVectors θP and θM.

get_hybridproblem_priors(::AbstractHybridProblem; scenario)

Return a dictionary of marginal prior distributions for components in θP and θM. Defaults for each component θ to Normal(θ, max(θ, 1.0)).

get_hybridproblem_train_dataloader([rng,] ::AbstractHybridProblem; scenario, n_batch)

Return a DataLoader that provides a tuple of

• xM: matrix of covariates, with one column per site
• xP: Iterator of process-model drivers, with one element per site
• y_o: matrix of observations with added noise, with one column per site
• y_unc: matrix sizeof(y_o) of uncertainty information

get_hybridproblem_transforms(::AbstractHybridProblem; scenario)

Return a NamedTupe of

• transP: Bijectors.Transform for the global PBM parameters, θP
• transM: Bijectors.Transform for the single-site PBM parameters, θM

get_hybridproblem_ϕunc(::AbstractHybridProblem; scenario)

Provide a ComponentArray of the initial additional parameters of the approximate posterior.

Create a loss function for parameter vector ϕ, given

• g(x, ϕ): machine learning model
• transPMS: transformation from unconstrained space to parameter space
• f(θMs, θP): mechanistic model
• interpreters: assigning structure to pure vectors, see negelbotransnorm_gf
• n_MC: number of Monte-Carlo sample to approximate the expected value across distribution

The loss function takes in addition to ϕ, data that changes with minibatch

• rng: random generator
• xM: matrix of covariates, sites in columns
• xP: drivers for the processmodel: Iterator of size n_site
• yo, yunc: matrix of observations and uncertainties, sites in columns

Create a loss function for parameter vector p, given

• g(x, ϕ): machine learning model
• f(θMs, θP): mechanistic model
• xM: matrix of covariates, sites in columns
• y_o: matrix of observations, sites in columns
• int_ϕθP: interpreter attachin axis with compponents ϕg and pc.θP

composition f ∘ transM ∘ g: mechanistic model after machine learning parameter prediction

(app::AbstractGPUDataHandler)(x) = handle_GPU_data(app, x)

Callable applied to argument x, used to configure the exchange of data between GPU and CPU. By Default, nothing is done to x. Package extensions for Flux and Lux implement overloads for AbstractGPUArray that call cpu(x) to transfer data on the GPU to CPU. Those package extension also use set_default_GPUHandler() so that .

init_hybrid_params(θP, θM, ϕg, n_batch; transP=asℝ, transM=asℝ)

Setup ComponentVector of parameters to optimize, and associated tools. Returns a NamedTuple of

• ϕ: A ComponentVector of parameters to optimize
• transPMsbatch, interpreters: Transformations and interpreters as required by `negelbotransnormgf`.
• gettransPMs: a function returning transformations `(nsite) -> (;P,Ms)`
• getcaintPMs: a function returning ComponentArrayInterpreter for PMs vector with PMs shaped as a matrix of `nsitecolumns ofθM`

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 with format of init_hybrid_ϕunc.

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.

Inverse of s = sumn(n) for positive integer n.

Gives an inexact error, if given s was not such a sum.

Cost function (ELBO) for hybrid model with batched sites.

It generates n_MC samples for each site, and uses these to compute the expected value of the likelihood of observations.

Arguments

• rng: random number generator (ignored on CUDA, if ϕ is a AbstractGPUArray)
• ϕ: flat vector of parameters interpreted by interpreters.μPϕgunc and interpreters.PMs
• g: machine learning model
• transPMs: Transformations as generated by gettransPMs returned from inithybrid_params
• f: mechanistic model
• py: negative log-likelihood of observations given predictions: function(y_ob, y_pred, y_unc)
• xM: matrix of covariates (ncov x nsite_batch)
• xP: model drivers, iterable of (nsitebatch)
• y_ob: matrix of observations (nobs x nsite_batch)
• y_unc: observation uncertainty provided to py (same size as y_ob)
• interpreters: NamedTuple as generated by gen_hybridcase_synthetic with entries:
  • μP_ϕg_unc: extract components of parameter of
    1. means of global PBM, 2) ML-weights, and 3) additional parameters of approximation q
  • PMs: assign components to PBM parameters 1 global, 2 matrix of n_site column vectors
  • int_unc (can be omitted, if μP_ϕg_unc(ϕ).unc is already a ComponentVector)
• n_MC: number of MonteCarlo samples from the distribution of parameters to simulate using the mechanistic model f.

neg_logden_indep_normal(obs, μ, logσ2s; σfac=1.0)

Compute the negative Log-density of θM for multiple independent normal distributions, given estimated means μ and estimated log of variance parameters logσ2s.

All the arguments should be vectors of the same length. If obs, μ are given as a matrix of several column-vectors, their summed Likelihood is computed, assuming each column having the same logσ2s.

Keyword argument σfac can be increased to put more weight on achieving a low uncertainty estimate and means closer to the observations to help an initial fit. The obtained parameters then can be used as starting values for a the proper fit with σfac=1.0.

predict_gf(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix, interpreters;
    get_transPMs, get_ca_int_PMs, n_sample_pred=200, 
    gpu_data_handler=get_default_GPUHandler())

Prediction function for hybrid model. Returns an NamedTuple with entries

• θ: ComponentArray (n_θP + n_site * n_θM), n_sample_pred) of PBM model parameters.
• y: Array (n_obs, n_site, n_sample_pred) of model predictions.

Compute predictions and log-Determinant of the transformation at given transformed parameters for each site.

The number of sites is given by the number of columns in Ms, which is determined by the transformation, transPMs.

Steps:

• transform the parameters to original constrained space
• Applies the mechanistic model for each site

Extract relevant parameters from θ and return n_MC generated draws together with the vector of standard deviations, σ.

Arguments

int_unc: Interpret vector as ComponentVector with components ρsP, ρsM, logσ2logP, coeflogσ2_logMs(intercept + slope),

scale_centered_at(x, m, σrel=1.0)
scale_centered_at(x, m, σ)

Centers and rescales rows of matrix x around vector m. The scale can either be given relative to m or specified as a vector of same size as m.

select_ml_engine(;scenario)

Returns a value type Val{:Symbol} to dispatch on the machine learning engine to use.

• defaults to Val(:SimpleChains)
• :use_Lux ∈ scenario -> Val(:Lux)
• :use_Flux ∈ scenario -> Val(:Flux)

transformU_block_cholesky1(v::AbstractVector, cor_ends)

Transform a parameterization v of a blockdiagonal of upper triangular matrices into the this matrix. cor_ends is an AbstractVector of Integers specifying the last column of each block. 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.

Takes a vector of entries of a lower UnitUpperTriangular matrix and transforms it to an UpperTriangular that satisfies diag(U' * U) = 1.

This can be used to fit parameters that yield an upper Cholesky-Factor of a Covariance matrix.

It uses the upper triangular matrix rather than the lower because it involes a sum across columns, whereas the alternative of a lower triangular uses sum across rows. Sum across columns is often faster, because entries of columns are contiguous.

Extract entries of upper diagonal matrix of UppterTriangular to columnwise vector

Compute the index in the vector of entries in an upper tridiagonal matrix

Extract entries of upper diagonal matrix of UnitUppterTriangular to columnwise vector

Convert vector v columnwise entries of upper diagonal matrix to UppterTriangular

Compute the (one-based) position (row, col) within an upper tridiagonal matrix for given (one-based) position, s within a packed vector representation.