Reference

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.

AbstractHVIApproximation

Provides a type hierarchy to distinguish different forms and parameterizations of posterior approximations.

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_ϕq
• get_hybridproblem_train_dataloader (may use construct_dataloader_from_synthetic)
• get_hybridproblem_test_data
• 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
• ϕq: additional parameters of the approximte posterior: returned by get_hybridproblem_ϕq

AbstractModelApplicator(x, ϕ; is_testmode = false)

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 overloads of construct_ChainsApplicator for layers of

• SimpleChains.SimpleChain
• Flux.Chain
• Lux.Chain

Abstraction of applying a process-based model with global parameters, θP, site-specific parameters, θMs (sites in columns), and site-specific model drivers, xP (sites in columns), It returns a matrix of predictions sites in columns.

Specific implementations need to provide function apply_model(app, θP, θMs_tr, xP). where

• θsP and θsMs_tr are shaped according to the output of generate_ζ, i.e. (n_site_pred x n_par x n_MC). Note that this transposed shape is different from most other parts of the interface, where sites are in the last dimension. The reason is that a column of a parameter is more efficient to transform between constrain and unconstrained scale.
• Results are of shape (n_obs x n_site_pred x n_MC).

They may also provide function apply_model(app, θP, θMs_tr, xP) for a sample of parameters, i.e. where an additional dimension is added to both θP and θMs. However, there is a default implementation that mapreduces across these dimensions.

Provided are implementations

• PBMSiteApplicator: based on a function that computes predictions per site
• PBMPopulationApplicator: based on a function that computes predictions for entire population
• NullPBMApplicator: returning its input θMs for testing
• DirectPBMApplicator: based on a function that takes the same arguments as apply_model

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)
ComponentArrayInterpreter(n_dims::NTuple{N,<:Integer}, ::AbstractComponentArray, m_dims::NTuple{M,<:Integer})

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]

'''

DirectPBMApplicator()

Process-based-Model applicator that invokes directly given function f(θP::AbstractVector, θMs::AbstractMatrix, xP).

Exp()

A bijector that applies broadcasted exponential function, i.e. exp.(x). It is equivalent to elementwise(exp) but works better with automatic differentiation on GPU.

Implements AbstractHybridProblem by gathering all the parts into one struct.

Fields:

• θP::ComponentVector, θM::ComponentVector: parameter templates
• g::AbstractModelApplicator, ϕg::AbstractVector: ML model and its parameters
• ϕq::ComponentVector: parameters for the Covariance matrix of the approximate posterior
• f_batch: Process-based model predicting for n_batch sites
• priors: AbstractDict: Prior distributions for all PBM parameters on constrained scale
• py: Likelihood function
• transM::Stacked, transP::Stacked: bijectors transforming from unconstrained to constrained scale for site-specific and global parameters respectively.
• train_dataloader::MLUtils.DataLoader: providing Tuple of matrices (xM, xP, y_o, y_unc, i_sites): covariates, model drivers, observations, observation uncertainties and index of provided sites.
• test_data::Tuple of the same form as withtrain_dataloader` for testset data.
• n_covar::Int, n_site::Int, n_batch::Int: number covariates, number of sites, and number of sites within one batch
• cor_ends::NamedTuple: block structure in correlations, defaults to (P = [length(θP)], M = [length(θM)])
• pbm_covars::NTuple{N,Symbol}: names of global parameters used as covariates in the ML model, defaults to (), i.e. no covariates fed into the ML model
• approx::AbstractHVIApproximation: the approximation to use for the variational distribution, defaults to MeanHVIApproximationMat(), i.e. correlation matrix shared across all sites and variance parameters a function of the predicted parameters.

HybridProblem(prob::AbstractHybridProblem; scenario = ()

Gather all information from another AbstractHybridProblem with possible updating of some of the entries.

Logistic()

A bijector that applies broadcasted exponential function, i.e. logit.(x). It is equivalent to elementwise(exp) but works better with automatic differentiation on GPU.

MagnitudeModelApplicator(app, y0)

Wrapper around AbstractModelApplicator that multiplies the prediction of the wrapped app by scalar y0.

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.

NormalScalingModelApplicator(app, lowers, uppers, FT::Type; repeat_inner::Integer = 1)

Fit a Normal distribution to number iterators lower and upper and transform results of the wrapped app AbstractModelApplicator. If repeat_inner is given, each fitted distribution is repeated as many times to support independent multivariate normal distribution.

FT is the specific FloatType to use to construct Distributions, It usually corresponds to the type used in other ML-parts of the model, e.g. Float32.

NullModelApplicator()

Model applicator that returns its inputs. Used for testing.

NullPBMApplicator()

Process-Base-Model applicator that returns its θMs inputs. Used for testing.

PBMPopulationApplicator(fθpop, n_site; θP, θM, θFix, xPvec)

Construct AbstractPBMApplicator from process-based model fθ that computes predictions across sites for a population of size n_site. The applicator combines enclosed θFix, with provided θMs and θP to a ComponentMatrix with parameters with one row for each site, that can be column-indexed by Symbols.

Arguments

• fθpop: process model, process model f(θc, xPc), which is agnostic of the partitioning of parameters into fixed, global, and individual.
  • θc: parameters: ComponentMatrix (nsite x npar) with each row a parameter vector
  • xPc: observations: ComponentMatrix (nobs x nsite) with each column
  observationsfor one site
• n_site: number of indiduals, i.e. rows in θMs
• θP: ComponentVector template of global process model parameters
• θM: ComponentVector template of individual process model parameters
• θFix: ComponentVector of actual fixed process model parameters
• xPvec: ComponentVector template of model drivers for a single site

PBMPopulationGlobalApplicator(fθpop, n_site; θP, θM, θFix, xPvec)

Construct AbstractPBMApplicator from process-based model fθ that computes predictions across sites for a population of size n_site. The applicator combines enclosed θFix, with provided θMs and θP to a ComponentMatrix with parameters with one row for each site, that can be column-indexed by Symbols.

Arguments

• fθpop: process model, process model f(θsc_tr, θgc, xPc), which is not agnostic of the partitioning of parameters into fixed, global, and individual to increase performance
  • θsc_tr: parameters: ComponentMatrix (nsite x npar_site) with each row a parameter vector
  • θgc: parameters: ComponentVector (nparglobal)
  • xPc: observations: ComponentMatrix (nobs x nsite) with each column
  observationsfor one site
• n_site: number of indiduals, i.e. rows in θMs
• θP: ComponentVector template of global process model parameters
• θM: ComponentVector template of individual process model parameters
• θFix: ComponentVector of actual fixed process model parameters
• xPvec: ComponentVector template of model drivers for a single site

PBMSiteApplicator(fθ; θP, θM, θFix, xPvec)

Construct AbstractPBMApplicator from process-based model fθ that computes predictions for a single site. The Applicator combines enclosed θFix, with provided θMs and θP and constructs a ComponentVector that can be indexed by symbolic parameter names, corresponding to the templates provided during construction of the applicator.

Arguments

• fθ: process model, process model fθ(θc, xP), which is agnostic of the partitioning

of parameters.

• θP: ComponentVector template of global process model parameters
• θM: ComponentVector template of individual process model parameters
• θFix: ComponentVector of actual fixed process model parameters
• xPvec:ComponentVector template of model drivers for a single site

RRuleMonitor(label, f, [ad_backend::ADTypes.AbstractADType])

Identity mapping of Callable or function f that intervenes the the pullback and raises an error if the supplied cotangent or the jacobian contains non-finitie entries.

Arguments

• label: id (String, or symbol) used in the error message.
• ad_backend: the AD backend used in DifferentiationInterface.jacobian. Defaults to AutoZygote().

RangeScalingModelApplicator(app, y0)

Wrapper around AbstractModelApplicator assumed to predict on (0,1) with a linear mappting to prededfined range.

RangeScalingModelApplicator(app, lowers, uppers, FT::Type; repeat_inner::Integer = 1)

Provide the target ragen by vectors lower and upper. The size of both outputs must correspond to the size of the output of app.

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.

apply_preserve_axes(f, ca::ComponentArray)

Apply callable f(x) to the data inside ca, assume that the result has the same shape, and return a new ComponentArray with the same axes as in ca.

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::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.

cpu_ca(ca::CA.ComponentArray)

Move ComponentArray form gpu to cpu.

creatensiteapplicator(app::AbstractPBMApplicator, n_site)

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.

gdev_hybridproblem_dataloader(dataloader::MLUtils.DataLoader; gdev_M, gdev_P,
    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; n_site_test = 0, 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
• y_true: matrix of site-specific observations with one column per site
• y_o: observations with added noise

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. For an empty vector Int[0] is returned.

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_gcdev(scenario::Val{scen}) where scen

Configure the function that puts data and computations to gpu device for given scenario. Checking for :use_gpu and :f_on_gpu in scenario. Returns a NamedTuple (;gdev_M, gdev_P)

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_correlation_Ms(prob; scenario = Val(()))

Extract correlation matrix of a problem based on MeanHVIApproximation. At unconstrained parameter scale.

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 a template tuple of ComponentVectors θP and θM at unconstrained (original) scale.

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)) from parameter templates.

get_hybridproblem_train_dataloader(::AbstractHybridProblem; scenario, n_batch)
get_hybridproblem_test_data end::AbstractHybridProblem; scenario)

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 the minibatch

get_hybridproblem_transforms(::AbstractHybridProblem; scenario)

Return a NamedTuple of

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

get_hybridproblem_θP(::AbstractHybridProblem; scenario)

Provide current expectations of population level PBM parameters on constrained original scale. Defaults to inverse-transform of ϕq.μP.

get_hybridproblem_ϕq(::AbstractHybridProblem; scenario)

Provide a ComponentArray of the non-ML parameters. Usually those contain the means of the unconstrained scale population parameters, μP.

Create a loss function for given

• g(x, ϕ): machine learning model
• transM: transformation of parameters at unconstrained space
• f(θMs_tr, θP): mechanistic model
• py: function(y_pred, y_obs, y_unc) to compute negative log-likelihood, i.e. cost
• intϕ: interpreter attaching axis with components ϕg and ϕP
• 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
• yunc: vector of uncertainty information for each observation Currently, hardcoes squared error loss of `(ypred .- yo) ./ σ, withσ = exp.(yunc ./ 2)`.
• i_sites: index of sites in the batch

and returns a NamedTuple of

• nLjoint: the negative-log of the joint parameter probability (Likelihood * prior)
• y_pred: predicted values
• θMs_tr, θP: PBM-parameters
• nLy: negative log-Likelihood of y_pred
• neg_log_prior: negative log-prior of θMs_tr and θP
• neg_log_prior: negative log-prior of θMs_tr and θP

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

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_ϕq(θP, θM, transP; cor_ends, n_site)

Initialize the non-ML parameter vector.

init_hybrid_ϕunc(approx::AbstractHVIApproximation, cor_ends, ρ0=0f0; 
  logσ2_ζP, coef_logσ2_ζMs, ρsP, ρsM)

Initialize vector of additional parameter of the approximate posterior.

Arguments:

• approx: AbstractMeanHVIApproximation, which is used. Parametrization will differ depending on the approximation.
• 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.

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.ϕg_ϕq 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:
  • ϕg_ϕq: 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_ϕq interpreter of components of ϕq
• 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 obs 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.

ones_similar_x(x::AbstractArray, size_ret = size(x))

Return ones(eltype(x), size_ret). Overload this methods for specific AbstractGPUArrays to return the correct container type. See e.g. HybridVariationalInferenceCUDAExt that calls CUDA.fill to return a CuArray rather than Array.

predict_hvi([rng], prob::AbstractHybridProblem)

Prediction function for hybrid variational inference parameter model.

Arguments

• prob: The problem for which to predict

Keyword arguments

• scenario
• n_sample_pred
• xM: covariates for the machine-learning model (ML): Matrix (nθM x nsitepred). Possibility to override the default from `gethybridproblemtraindataloader`.
• xP: model drivers for process based model (PBM): Matrix with (nsitepred) rows. Possibility to override the default from get_hybridproblem_train_dataloader.

Returns an NamedTuple (; y, θsP, θsMs_tr, 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_tr: 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.

Note that for some approximations, such as MeanVarSepHVIApproximation, prob.ϕq contains uncertainty parameters that are specific to sites. For a proper prediction for new sites, update prob.ϕq before appropriately. If xM and xP is unspecified, the problem's sites are used and prob.ϕq is consistent. If predicting for a subset of sites, specify keyword argument i_sites .

predict_point_hvi([rng], prob::AbstractHybridProblem)

Prediction function for hybrid variational inference parameter model that omits the sampling step but returns the prediction at the mean in unconstrained space.

Arguments

• prob: The problem for which to predict

Keyword arguments

• scenario
• gdevs
• xM: covariates for the machine-learning model (ML): Matrix (nθM x nsitepred). Possibility to override the default from `gethybridproblemtraindataloader`.
• xP: model drivers for process based model (PBM): Matrix with (nsitepred) rows. Possibility to override the default from get_hybridproblem_train_dataloader.

Returns an NamedTuple (; y, θMs_tr, θP) with entries

• y: Matrix (n_obs, n_site) of model predictions.
• θP: ComponentVector of PBM model parameters that are kept constant across sites.
• θMs_tr: ComponentMatrix (n_site, n_θM) of PBM model parameters that vary by site.

sample_posterior(rng, prob, [xM::AbstractMatrix]; scenario=Val(()), kwargs...)

Sampling the posterior parameter distribution for hybrid variational inference problem.

Arguments

• rng: random number generator
• prob: The AbstractHybridProblem from to sample
• xM: covariates for the machine-learning model (ML): Matrix (n_θM x n_site_pred). Default to all sites in get_hybridproblem_train_dataloader(prob; scenario).

Optional keyword arguments

• scenario: scenario to query prob and set default of gpu devices.
• n_sample_pred: number of samples to draw, defaults to 200
• gdevs: NamedTuple(gdev_M, gdev_P): GPU devices for machine learning model and parameter transformtation, default to get_gdev_MP(scenario).
• is_inferred: set to Val(true) to activate type stabilicy check for transformation

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

• θsP: ComponentArray (n_θP, n_sample_pred) of PBM model parameters that are kept constant across sites.
• θsMs_tr: 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.

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)).

zero_penalty_loss(y_pred, θMs, θP, ϕg, ϕq)

Add zero i.e. no additional loss terms during the HVI fit.

The basic cost in HVI is the negative log of the joint probability, i.e. the likelihood of the observations given the parameters * prior probability of the parameters.

Sometimes there is additional knowledge not encoded in the prior, such as one parameter must be larger than another, or entropy-weights of the ML-parameters, and the solver accept a function to add additional loss terms.

Arguments

• y_pred::AbstractMatrix: Observations
• θMs::AbstractMatrix: site parameters
• θP::AbstractVector: global parameters
• ϕg: ML-model parameters,
• ϕq::AbstractVector, additional parameters of the posterior

f_doubleMM(θc_tr::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_tr 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

function f_doubleMM(θc_tr::ComponentVector{ET}, x) where ET
    # extract parameters not depending on order, i.e whether they are in θP or θM
    # r0 = θc_tr[:r0]
    (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
        getdata(θc_tr[par])::ET
    end
    y = r0 .+ r1 .* x.S1 ./ (K1 .+ x.S1) .* x.S2 ./ (K2 .+ x.S2)
    return (y)
end

f_doubleMM_sites(θc_tr::CA.ComponentMatrix, xPc::CA.ComponentMatrix)

Example process based model (PBM) that predicts for a batch of sites.

Arguments

• θc_tr: 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.

function f_doubleMM_sites(θc_tr::ComponentMatrix, xPc::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 vectors that are row-repeated into a matrix
    VT = typeof(CA.getdata(θc_tr)[:,1])   # workaround for non-type-stable Symbol-indexing
    n_obs = size(S1, 1)
    rep_fac = ones_similar_x(xPc, n_obs)      # to reshape into matrix, avoiding repeat
    (r0, r1, K1, K2) = map((:r0, :r1, :K1, :K2)) do par
        p1 = CA.getdata(θc_tr[:, par]) ::VT
        #(r0 .* rep_fac)'    # move to computation below to save allocation
        #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)
    (r0 .* rep_fac)' .+ (r1 .* rep_fac)' .* S1 ./ ((K1 .* rep_fac)' .+ S1) .* S2 ./ ((K2 .* rep_fac)' .+ S2)
end