HybridVariationalInference

AbstractComponentArrayInterpreter

Interface for Type that implements

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

When called on a vector, forwards to as_ca.

There is a default implementation for Base.length based on ComponentArrays.getaxes.

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 (may use construct_dataloader_from_synthetic)
• get_hybridproblem_priors
• get_hybridproblem_n_covar
• get_hybridproblem_n_site_and_batch

optionally

• gen_hybridproblem_synthetic
• get_hybridproblem_float_type (defaults to eltype(θM))
• get_hybridproblem_cor_ends (defaults to include all correlations: (P = [length(θP)], M = [length(θM)]))
• get_hybridproblem_pbmpar_covars (defaults to empty tuple)

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]

'''

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], such 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.

Fit a Normal distribution to iterators lower and upper. If repeat_inner is given, each fitted distribution is repeated as many times.

NullModelApplicator()

Model applicator that returns its inputs. Used for testing.

StackedArray(stacked, nrow)

A Bijectors.Transform that applies stacked to each column of an n-row matrix.

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.

Compute predictions of the transformation at given transformed parameters for each site. The number of sites is given by the number of rows in ζsMs.

Steps:

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

ζsP and ζsMs are shaped according to the output of generate_ζ. Results are of shape (n_obs x n_site_pred x n_MC).

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

compose_axes(axtuples::NamedTuple)

Create a new 1d-axis that combines several other named axes-tuples such as of key = getaxes(::AbstractComponentArray).

The new axis consists of several ViewAxes. If an axis-tuple consists only of one axis, it is used for the view. Otherwise a ShapedAxis is created with the axes-length of the others, essentially dropping component information that might be present in the dimensions.

Compute the cholesky-factor parameter for a given single correlation in a 2x2 matrix. Invert the transformation of cholesky-factor parameterization.

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.

Compute the components of the elbo for given initial conditions of the problems for the first batch of the trainloader.

construct_3layer_MLApplicator(
    rng::AbstractRNG, prob::HVI.AbstractHybridProblem, <ml_engine>;
    scenario::Val{scen}) where scen

Construct a machine learning model for given Proglem and machine learning engine. Implemented for machine learning extensions, such as Flux or SimpleChains. ml_engine usually is of type Val{Symbol}, e.g. Val(:Flux). See select_ml_engine.

Scenario is a value-type of NTuple{_,Symbol}.

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

construct_dataloader_from_synthetic(rng::AbstractRNG, prob::AbstractHybridProblem;
    scenario = (), n_batch)

Construct a dataloader based on gen_hybridproblem_synthetic.

Construct a parametric type-stable model applicator, given covariates, x, and parameters, ϕ.

The default returns the current model applicator.

In order to let mean of θ stay close to initial point parameter estimates construct a prior on mean θ to a Normal around initial prediction.

cpu_ca(ca::CA.ComponentArray)

Move ComponentArray form gpu to cpu.

extend_stacked_nrow(b::Stacked, nrow::Integer)

Create a Stacked bijectors that transforms nrow times the elements of the original Stacked bijector.

Example

X = reduce(hcat, ([x + y for x in 0:4 ] for y in 0:10:30))
b1 = CP.Exp()
b2 = identity
b = Stacked((b1,b2), (1:1,2:4))
bs = extend_stacked_nrow(b, size(X,1))
Xt = reshape(bs(vec(X)), size(X))
@test Xt[:,1] == b1(X[:,1])
@test Xt[:,2:4] == b2(X[:,2:4])

flatten1(cv::CA.ComponentVector)

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

gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader,
    scenario = (), 
    gdev = gpu_device(),
    gdev_M = :use_gpu ∈ scenario ? gdev : identity,
    gdev_P = :f_on_gpu ∈ scenario ? gdev : identity,
    batchsize = dataloader.batchsize,
    partial = dataloader.partial
    )

Put relevant parts of the DataLoader to gpu, depending on scenario.

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_hybridproblem_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_ζresid_norm to the means extracted from parameters and predicted by the machine learning model.

The output shape of size (n_site x n_par x n_MC) is tailored to iterating each MC sample and then transforming each parameter on block across sites.

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.

get_hybridproblem_n_site_and_batch(::AbstractHybridProblem; scenario)

Provide the number of sites.

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(::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
• i_sites: Vector of indices of sites in toal sitevector for the minibatch

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. Defaults to zero correlation and log_σ2 of 1e-10.

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 neg_elbo_gtf
• n_MC: number of Monte-Carlo sample to approximate the expected value across distribution
• pbm_covars: tuple of symbols of process-based parameters provided to the ML model
• θP: ComponentVector as a template to select indices of pbm_covars

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
• y_o, y_unc: matrix of observations and uncertainties, sites in columns

Create a loss function for given

• g(x, ϕ): machine learning model
• transM: transforamtion of parameters at unconstrained space
• f(θMs, θP): mechanistic model
• yoglobal: site-independent observations
• 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

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
• i_sites: index of sites in the batch

get_positions(cai::AbstractComponentArrayInterpreter)

Create a NamedTuple of integer indices for each component. Assumes that interpreter results in a one-dimensional array, i.e. in a ComponentVector.

Get the inverse-transformation of lower and upper quantiles of a Vector of Distributions.

This can be used to get proper confidence intervals at unconstrained (log) ζ-scale for priors on normal θ-scale for constructing a NormalScalingModelApplicator.

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

composition transM ∘ g: transformation after machine learning parameter prediction Provide a transMs = StackedArray(transM, n_batch)

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 `negelbo_gtf`.
• 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_ζP, coef_logσ2_ζMs, ρ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_ζMs, defaults to [-10.0, 0.0]

Returns a ComponentVector of

• logσ2_ζP: vector of log-variances of ζP (on log scale). defaults to -10
• coef_logσ2_ζMs: 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.

Map process base model (PBM), f, across each site.

Arguments

• f(θ, xP, args...; intθ1, kwargs...): Process based model for single siteMake sure to hint the type, so that results can be inferred.
• θMst: transposed model parameters across sites matrix: (nparM, nsite_batch)
• θP: transposed model parameters that do not differ by site: (n_parP,)
• θFix: Further parameter required by f that are not calibrated.
• xP: Model drivers: Matrix with nsitebatch columns. If provided a ComponentArray with labeled rows, f can then access xP[:key].
• intθ1: ComponentArrayInterpreter that can be applied to θ, so that entries can be extracted.

See test_HybridProblem of using this function to construct a PBM function that can predict across all sites.

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, xP, y_ob, y_unc, i_sites: information of the sites in the current minibatch
  • 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)
  • i_sites: indices of sites for current minibatch
• interpreters: NamedTuple as generated by gen_hybridproblem_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.

Compute the neg_elbo for each sampled parameter vector (last dimension of ζs).

• Transform and compute log-jac
• call forward model
• compute log-density of predictions
• compute entropy of transformation

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_hvi([rng], prob::AbstractHybridProblem [,xM, xP]; scenario, ...)
predict_hvi(rng, g, f, ϕ::AbstractVector, xM::AbstractMatrix;
    get_transPMs, get_ca_int_PMs, n_sample_pred=200, gdev = identity)

Prediction function for hybrid variational inference parameter model.

Arguments

• The problem for which to predict
• xM: covariates for the machine-learning model (ML): Matrix (nθM x nsite_pred).
• xP: model drivers for process based model (PBM): Matrix with (nsitepred) rows. If provided a ComponentArray with a Tuple-Axis in rows, the PBM model can access parts of it, e.g. xP[:S1,...].

Keyword arguments

• scenario
• nsamplepred

Returns an NamedTuple (; y, θsP, θsMs, entropy_ζ) with entries

• y: Array (n_obs, n_site, n_sample_pred) of model predictions.
• θsP: ComponentArray (n_θP, n_sample_pred) of PBM model parameters that are kept constant across sites.
• θsMs: ComponentArray (n_site, n_θM, n_sample_pred) of PBM model parameters that vary by site.
• entropy_ζ: The entropy of the log-determinant of the transformation of the set of model parameters, which is involved in uncertainty quantification.

Extract relevant parameters from ζ and return nMC generated multivariate normal draws together with the vector of standard deviations, σ: `(ζPresids, ζMsparfirstresids, σ)The output shape(nθ, nsite?, nMC)is tailored to addingζMsparfirstresidsto ML-model predcitions of size(nθM, n_site)`.

Arguments

• int_unc: Interpret vector as ComponentVector with components ρsP, ρsM, logσ2ζP, coeflogσ2_ζMs(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)

stack_ca_int(cai::AbstractComponentArrayInterpreter, ::Val{n_dims})

Interpret the first dimension of an Array as a ComponentArray. Provide the Tuple of following dimensions by a value type, e.g. Val((n_col, n_z)).

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.

Transform parameters

• from unconstrained (e.g. log) ζ scale of format ((nsite x npar) x n_mc)
• to constrained θ scale of the same format

Transforms each row of a matrix (nMC x nPar) with site parameters Ms inside nPar of form (npar x nsite) to Ms of the form (nsite x n_par), i.e. neighboring entries (inside a column) are of the same parameter.

This format of having n_par as the last dimension helps transforming parameters on block.

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.