Reference of internal functions

OneBasedVectorWithZero(data)

A thin wrapper over an AbstractVector that exposes a linear 1-based indexing API mapping v[i] to data[axes(data, 1)[i]] on the underlying storage and provides a value at index 0 (defaulting to zero) that is not stored in the underlying vector. The zero is not counted in the length of the vector, so length(v) is the same as length(data).

Example usage:

v = HybridVariationalInference.OneBasedVectorWithZero([10,20,30])
v[1] == 10
v[2] == 20
v[3] == 30
v[0] == 0 # default value at index 0 is zero
v[[1,0,0,3]] == [10,0,0,30]

WeightedObsView{D}

A view over a dataset that supports weighted resampling via MLUtils.getobs.

solve(prob::AbstractHybridProblem, solver::HybridPosteriorSolver; epochs, ...)

Perform the inversion of HVI Problem.

Arguments

• prob: The AbstractHybridProblem to solve.
• scenario: Scenario to query prob, defaults to Val(()).
• epochs: number of epochs to train, i.e. number of passes through the whole dataset.

Optional keyword arguments

• rng: Random generator, defaults to Random.default_rng().
• gdevs: NamedTuple (;gdev_M, gdev_P) functions to move computation and data of ML model on and PBM respectively to gpu (e.g. gpu_device() or cpu (identity). defaults to get_gdev_MP(scenario)
• θmean_quant default to 0.0: deprecated
• is_inferred: set to Val(true) to activate type stability checks
• is_omit_priors: set to Val(true) to omit priors in the loss computation, which can be useful for debugging or if priors are not implemented for a specific scenario (e.g. on gpu)
• clusters: vector of cluster assignments for each site, defaults to each site being its own cluster. Clusters are used to compute the loss in a way that accounts for clustering of sites, which can be useful if there are many sites and the number of Monte Carlo samples is limited.
• cluster_rep: number of times to repeat each cluster in the loss computation, defaults to 1. Repeating clusters can be useful to effectively increase the number of Monte Carlo samples when the number of clusters is small.
• epochs_callback: number of epochs between progress output on evaluating testdata

Returns a NamedTuple of

• probo: A copy of the HybridProblem, with updated optimized parameters
• interpreters: TODO
• ϕ: the optimized HVI parameters: a ComponentVector with entries
  • ϕg: The ML model parameter vector,
  • ϕq: ComponentVector of non-ML parameters, including μP: ComponentVector of the mean global PBM parameters at unconstrained scale
• θP: ComponentVector of the mean global PBM parameters at constrained scale
• resopt: the structure returned by Optimization.solve. It can contain more information on convergence.

Create Cholesky-factor of Covariance matrix. from Cholesky-factors of correlation matrices (UP, UM) and diagonal sqrt(variances) (σP, σMs) Return this factor and its diagonal

as_ca(v::AbstractArray, interpretor)

Returns a ComponentArray with underlying data v.

overdispersion_test(Y, μ, Σ; α=0.05)

Test whether the q×p sample matrix Y (rows = individuals) is overdispersed relative to the reference distribution N(μ, Σ).

Returns: Sn, E0, Var0, Z, pvaluenormal, pvalue_chisq

colsum_finite_obs(X::AbstractMatrix, obs::AbstractMatrix)

Return the column sums of X, using obs as a finiteness mask. Equivalently, entries of X are treated as zero whenever the corresponding entry of obs is not finite.

Important: this function does not check whether entries of X are finite. If X[i, j] is NaN, Inf, or -Inf and obs[i, j] is finite, then that value is included in the sum.

X and obs must have the same axes. For ordinary Julia Matrix objects, this means they must have the same size.

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.

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

construct log-determinant of covariance matrix. which depends on the σ = sqrt(diag(cov)). The part for for θMs is reduced by considering that a minibatch contains only a fraction of the data.

compute_pvalue_asymptotic_overdispersion_from_dist2(dist2_matrix)

Compute p-value for overdispersion using asymptotic approximation, based on a precomputed matrix of squared Mahalanobis distances.

Arguments

• dist2_matrix: m × m symmetric matrix of squared Mahalanobis distances (dist2matrix[i,j] = (xi - xj)' Σ⁻¹ (xi - x_j))
• n: the dimension of x_i (number of variables)
• The matrix must be symmetric and contain only upper/lower triangle values

Returns

• p_value: one-sided p-value for overdispersion

Compute the elbo components for a single parameter sample. Site-specific vectors: nL, loss_penalty

omit the last n elements of an iterator

generate_repeated_integers(n_MC::Int, n_sample_ranef::Int) -> Vector{Int}

Generate a vector of increasing integers where each integer is repeated n_sample_ranef times, except possibly the last one, such that the total length of the vector is exactly n_MC.

Arguments

• n_MC::Int: The total length of the output vector. Must be a positive integer.
• n_sample_ranef::Int: The number of times each integer is repeated. Must be a positive integer.

Returns

• Vector{Int}: A vector of length n_MC where each integer i appears n_sample_ranef times, except for the last integer which appears mod(n_MC, n_sample_ranef) times if n_MC is not a multiple of n_sample_ranef, and n_sample_ranef times otherwise.

Examples

julia-repl julia> generate_repeated_integers(8, 5) 8-element Vector{Int64}: 1, 1, 1, 1, 1, 2, 2, 2

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_choleskyΣ_ranef(re::AbstractRandomEffectsComputer, ϕq_ranef::ComponentVector)

Return the upper cholesky factor of Covariance of the Random effects.

Create a loss function for parameter vector ϕ, given

• g(x, ϕ): machine learning model
• transPMS: transformation from unconstrained space to parameter space
• f(θMs_tr, θ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

Loss function for random effects, given fixed ζ.

get_zero_prior_logdensity(priorsP::Tuple, priorsM::Tuple, θP, θM)

invoke logpdf of prior and sum, to infer return type and proper zero of prior density.

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

insert_zeros(v, positions)

Return a new vector with zero(eltype(v)) inserted at each position in positions. Positions are applied in order against the growing vector (as if sequential inserts), so later indices are interpreted on the updated result. Only one output vector is allocated.

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

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

log_density_mvn_cholesky(U,x)

Compute the log-density of a zero-mean multivariate normal distribution with covariance matrix C = U' * U, where U is the upper Cholesky factor.

Arguments:

• U: upper triangular Cholesky factor of the covariance matrix (n × n)
• x: vector of length n (the sample)

Returns:

• log p(x) ∈ ℝ: log-density at x

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 joint density of predictions and unconstrained parameters, nLjoint and its components
  • nLy: The likelihood of the data, given the parameters
  • neg_log_prior: the prior of parameters at constrained scale
  • logjac, negative logarithm of the absolute value of the determinant of the Jacobian of the transformation θ=T(ζ).
• loss_penalty: additional loss terms from penalty_computer
• compute entropy of transformation

refit_clusters(rng, probo, solver, xM; scenario, n_cluster_initial, n_aggsplits, epochs)

Iteratively refit the model and split clusters of sites based on overdispersion tests.

When several sites are within one cluster, they are treated in a way, such that all the observations constrain the uncertainty of the mean estimate within that cluster. The fewer sites are within one cluster, the higher the uncertainty estimate.

This methods implements a strategy to start with few clusters and checks if the distribution of predicted site values within a cluster is overdispersed relative to the uncertainty predicted for the cluster. If so, the cluster is split into smaller clusters and the model is refitted. Because, the refitting changes the uncertainty estimates, only the few (1/10th) clusters with the most sites are checked for overdispersion, and then a refitting takes place before checking the next clusters.

Arguments

• rng: random number generator for reproducibility
• probo: the probabilistic model to fit
• solver: optimization algorithm for fitting
• xM: input data for the model
• scenario: optional argument for different scenarios in the model
• n_cluster_initial: number of clusters to start with
• n_aggsplits: number of clusters to split before refitting
• epochs: number of epochs for refitting after each series of splits

Returns

• probo: the refitted probabilistic model
• clusters: final cluster assignments for each site

replace_columns_matrix(x::Matrix, col_indices::Vector{Int}, y::Matrix)

Return a new matrix where the specified columns of x are replaced by the columns of y.

This function performs a non-mutating column replacement using one-hot projection matrices and is compatible with automatic differentiation frameworks like Zygote.

Parameters

• x: The input matrix of size (m, n) to be modified.
• col_indices: A vector of column indices (1-based) specifying which columns to replace.
• y: The replacement matrix of size (m, length(col_indices)).

Returns

• A new matrix of the same size as x, where the columns at positions col_indices are replaced by the corresponding columns of y.

Details

• The function constructs a one-hot projection matrix P_col (size n × length(col_indices)) using matrix comprehensions, where each column corresponds to a target column index.
• The replacement values are scattered into the full matrix space via x * (I - P_col * P_col') + y * P_col'.
• The operation is differentiable with respect to all inputs.

Example

x = [1 2 3; 4 5 6; 7 8 9]
col_indices = [1, 3]
y = [10 11; 12 13; 14 15]

result = replace_columns_matrix(x, col_indices, y)
#result = HVI.replace_columns_matrix(x, col_indices, y)
# result == [10 2 11; 12 5 13; 14 8 15]

Notes

• All column indices must be valid (1-based).
• The function is fully non-mutating and compatible with Zygote for automatic differentiation.
• The operation is differentiable with respect to all inputs.

replace_values_matrix(x::Matrix, i_sites::Vector{<:Integer}, pos::Vector{<:Integer}, y::Matrix)

Return a new matrix where the submatrix at positions x[i_sites, pos] is replaced by y, while all other values remain unchanged.

This function performs a non-mutating replacement using one-hot projection matrices and a binary mask derived from the outer product of indicator vectors. It is designed to be compatible with automatic differentiation frameworks like Zygote.

Parameters

• x: The input matrix of size (m, n) to be modified.
• i_sites: A vector of row indices (1-based) specifying which rows to replace.
• pos: A vector of column indices (1-based) specifying which columns to replace.
• y: The replacement matrix of size (length(i_sites), length(pos)).

Returns

• A new matrix of the same size as x, where x[i_sites, pos] is replaced by y.

Details

• The function constructs one-hot projection matrices P_row (size m × length(i_sites)) and P_col (size n × length(pos)) using matrix comprehensions.
• The replacement values are scattered into the full matrix space via P_row * y * P_col'.
• A binary mask is created using the outer product of indicator vectors for i_sites and pos.
• The result is computed as x .* (1 - mask) + (P_row * y * P_col') .* mask, blending the original matrix with the scattered replacement values.

Example

```julia x = [1 2 3; 4 5 6; 7 8 9] i_sites = [1, 3] pos = [2, 3] y = [10 11; 12 13]

result = replacevalues(x, isites, pos, y)

result = [1 10 11; 4 5 6; 7 12 13]

reshape_penalty_matrix(penalty::NamedTuple{KEYS}) where KEYS
reshape_penalty_matrix(penalty::ComponentVector{ET,KEYS}) where {ET,KEYS}

Reshape the output of the penalty computer to a ComponentMatrix. Assuming that all the component in penalty are of the same element type and length.

provide a sample npar x nsite x nsample, where npar is the dimension of parameters including those columns that have no random effect.

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

sum x across subarray for each entry in dim_ret. This is opposite of the dims argument in sum.

take_n!(itr, n)

Peel off the first n elements of an drop-iterator itr and return them as a vector, while mutating itr to now start after those n elements.

Examples

it = HybridVariationalInference.drop_iterate(1:5) # initialize the iterator

a1 = HybridVariationalInference.take_n!(it,3)
collect(a1) == [1,2,3]

a2 = HybridVariationalInference.take_n!(it,3)
collect(a2) == [4,5]  # only two element left, so return those

a3 = HybridVariationalInference.take_n!(it,3)
collect(a3) == [] # no elements left, so return empty vector

transformU_block_cholesky1(v::AbstractVector, cor_ends)

Transform a parameterization, v, of a blockdiagonal of upper triangular matrices into a vector with a matrix for each block. 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.

An correlation parameterization can parameterize a block of a single parameter, or an empty parameter block. To indicate the empty block, provide cor_ends == [0].

Takes a vector of parameters for 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 correlation matrix.

It uses the upper triangular matrix rather than the lower because it involves 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.

An empty parameterization

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.

create_ϕq(θP, ϕunc, transP::Stacked)

Add information on μP to ϕunc.

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.

vectuptotupvec(vectup)
vectuptotupvec_allowmissing(vectup)

Typesafe convert from Vector of Tuples to Tuple of Vectors. The first variant does not allow for missing in vectup. The second variant allows for missing but has eltype of Union{Missing, ...} in all components of the returned Tuple, also when there were not missing in vectup.

Arguments

• vectup: A Vector of identical Tuples

Examples

vectup = [(1,1.01, "string 1"), (2,2.02, "string 2")] 
HybridVariationalInference.vectuptotupvec_allowmissing(vectup) == 
  ([1, 2], [1.01, 2.02], ["string 1", "string 2"])

with_logabsdet_jacobians

Similar to withlogabsdetjacobian, but returns as the second component a vector of Jacobians of transformation of each component in x.

with_logabsdet_jacobians

Similar to withlogabsdetjacobian, but returns as the second component a vector of Jacobians of transformation of each component in x, rather than the sum.

with_logabsdet_jacobians(sb::StackedArray, x::AbstractArray)

Return a Jacobian for each row in x.

MLUtils.getobs(wov::WeightedObsView, indices)

Draw length(indices) weighted samples (with replacement) using StatsBase.sample with Weights, then fetch those observations from the underlying dataset. Both sampling and data fetching are shielded from Zygote via ChainRulesCore.@ignore_derivatives.

MLUtils.numobs(wov::WeightedObsView)

Return the number of observations in the underlying dataset.