romtools.workflows.inverse.full_covariance_vi#

Utilities for full-covariance Gaussian variational inference.

The full-covariance variational family is represented in optimizer coordinates as q(x) = N(mu, Sigma) with Sigma = L L.T. The Cholesky factor is a numerical representation; natural-gradient directions are expressed in the mean and covariance tangent spaces.

Functions

correlation_cholesky_from_covariance(covariance)

Return marginal standard deviations and a correlation Cholesky factor.

covariance_from_cholesky(cholesky)

Return the symmetric covariance represented by a lower Cholesky factor.

entropy(cholesky[, scaling_factor])

Return Gaussian entropy, optionally scaled with the ELBO.

estimate_natural_gradient(optimizer_samples, ...)

Estimate full-Gaussian natural and ordinary ELBO gradients.

infer_svec_dimension(size)

Infer the matrix dimension associated with an svec length.

log_density(optimizer_samples, mean, cholesky)

Evaluate a full-covariance Gaussian log density in optimizer coordinates.

marginal_std_from_cholesky(cholesky)

Return marginal standard deviations for a covariance Cholesky factor.

natural_score_terms(optimizer_samples, mean, ...)

Return per-sample natural scores for mean and covariance.

packed_direction_to_covariance(...)

Convert an svec covariance direction back to matrix form.

rescale_cholesky_marginals(cholesky, ...)

Rescale covariance marginals while preserving its correlation matrix.

retract_covariance(cholesky, ...[, ...])

Apply an SPD-preserving covariance retraction.

robust_cholesky(covariance[, ...])

Compute a Cholesky factor, adding only roundoff-scale jitter if needed.

svec(matrix)

Isometrically vectorize a symmetric matrix.

svec_inverse(vector[, dimensionality])

Invert svec().

svec_size(dimensionality)

Return the number of independent entries in a symmetric matrix.

romtools.workflows.inverse.full_covariance_vi.correlation_cholesky_from_covariance(covariance)[source]#

Return marginal standard deviations and a correlation Cholesky factor.

Parameters:

covariance (ndarray)

Return type:

Tuple[ndarray, ndarray]

romtools.workflows.inverse.full_covariance_vi.covariance_from_cholesky(cholesky)[source]#

Return the symmetric covariance represented by a lower Cholesky factor.

Parameters:

cholesky (ndarray)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.entropy(cholesky, scaling_factor=1.0)[source]#

Return Gaussian entropy, optionally scaled with the ELBO.

Parameters:
  • cholesky (ndarray)

  • scaling_factor (float)

Return type:

float

romtools.workflows.inverse.full_covariance_vi.estimate_natural_gradient(optimizer_samples, mean, cholesky, scaled_log_joint_terms, baseline_method, elbo_scaling_factor=1.0, entropy_strategy='analytic')[source]#

Estimate full-Gaussian natural and ordinary ELBO gradients.

scaled_log_joint_terms must already include the ELBO scaling factor and, for transformed bounded parameters, the transform-Jacobian contribution. Analytic entropy is added exactly after the stochastic log-joint estimate.

Parameters:
  • optimizer_samples (ndarray)

  • mean (ndarray)

  • cholesky (ndarray)

  • scaled_log_joint_terms (ndarray)

  • baseline_method (str)

  • elbo_scaling_factor (float)

  • entropy_strategy (str)

Return type:

dict

romtools.workflows.inverse.full_covariance_vi.infer_svec_dimension(size)[source]#

Infer the matrix dimension associated with an svec length.

Parameters:

size (int)

Return type:

int

romtools.workflows.inverse.full_covariance_vi.log_density(optimizer_samples, mean, cholesky)[source]#

Evaluate a full-covariance Gaussian log density in optimizer coordinates.

Parameters:
  • optimizer_samples (ndarray)

  • mean (ndarray)

  • cholesky (ndarray)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.marginal_std_from_cholesky(cholesky)[source]#

Return marginal standard deviations for a covariance Cholesky factor.

Parameters:

cholesky (ndarray)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.natural_score_terms(optimizer_samples, mean, cholesky)[source]#

Return per-sample natural scores for mean and covariance.

The covariance scores are returned in svec coordinates.

Parameters:
  • optimizer_samples (ndarray)

  • mean (ndarray)

  • cholesky (ndarray)

Return type:

Tuple[ndarray, ndarray]

romtools.workflows.inverse.full_covariance_vi.packed_direction_to_covariance(packed_direction, dimensionality)[source]#

Convert an svec covariance direction back to matrix form.

Parameters:
  • packed_direction (ndarray)

  • dimensionality (int)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.rescale_cholesky_marginals(cholesky, minimum_std, maximum_std)[source]#

Rescale covariance marginals while preserving its correlation matrix.

Parameters:
  • cholesky (ndarray)

  • minimum_std (ndarray | float)

  • maximum_std (ndarray | float)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.retract_covariance(cholesky, covariance_direction, step_size, max_covariance_log_step=1.0)[source]#

Apply an SPD-preserving covariance retraction.

The requested log-step is uniformly rescaled, rather than clipping individual eigenvalues, so the tangent direction is preserved.

Returns:

applied_scale multiplies step_size after log-step limiting.

Return type:

new_cholesky, applied_scale

Parameters:
  • cholesky (ndarray)

  • covariance_direction (ndarray)

  • step_size (float)

  • max_covariance_log_step (float | None)

romtools.workflows.inverse.full_covariance_vi.robust_cholesky(covariance, initial_jitter=0.0, max_jitter=1e-06)[source]#

Compute a Cholesky factor, adding only roundoff-scale jitter if needed.

Parameters:
  • covariance (ndarray)

  • initial_jitter (float)

  • max_jitter (float)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.svec(matrix)[source]#

Isometrically vectorize a symmetric matrix.

Lower-triangular entries are stored in NumPy’s tril_indices order. Off-diagonal entries are multiplied by sqrt(2) so Euclidean inner products of packed vectors equal Frobenius inner products of the matrices.

Parameters:

matrix (ndarray)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.svec_inverse(vector, dimensionality=None)[source]#

Invert svec().

Parameters:
  • vector (ndarray)

  • dimensionality (int | None)

Return type:

ndarray

romtools.workflows.inverse.full_covariance_vi.svec_size(dimensionality)[source]#

Return the number of independent entries in a symmetric matrix.

Parameters:

dimensionality (int)

Return type:

int