qpots.utils.tc_utils

Total-correlation, mutual-information, and multitask Gaussian-process helpers for partial-information optimization workflows.

qpots.utils.tc_utils._stable_logdet(A: ndarray, jitter: float = 1e-10, max_tries: int = 8) → float

Numerically stable log(det(A)) for (near) PSD matrices by adding diagonal jitter if needed. Raises if it cannot make the matrix numerically positive definite.

qpots.utils.tc_utils._unravel_index(flat_idx: int, shape)

Convert a flat index into a multidimensional index tuple.

Parameters:

• flat_idx (int) – Index into the flattened version of an array.

• shape (Sequence[int]) – Shape of the original array.

Returns:

Index tuple equivalent to numpy.unravel_index(flat_idx, shape).

Return type:

tuple[int, …]

qpots.utils.tc_utils.argmax_mi_subset_bruteforce(cov_or_samples: ndarray, *, subset_size: int | None = None, jitter: float = 1e-10, base: float = 2.718281828459045, assume_samples: bool | None = None, deduplicate_complements: bool = True, return_all_scores: bool = False) → Dict[str, Any]

Brute-force search over output subsets under a Gaussian assumption.

You can pass either:

• cov_or_samples as a (K,K) covariance matrix, OR

• cov_or_samples as (N,K) samples (rows=samples), from which we estimate covariance.

Parameters:

• cov_or_samples (np.ndarray) – (K,K) covariance OR (N,K) samples.

• subset_size (int or None) – If provided, restrict search to subsets whose size equals subset_size. If None, searches all non-trivial subsets (excluding empty/full).

• jitter (see above)

• base (see above)

• assume_samples (bool or None) – If None, auto-detect: (K,K) => covariance; otherwise => samples.

• deduplicate_complements (bool) – If subset_size is None, avoid evaluating both a subset and its complement by restricting the search to subsets of size at most floor(K / 2). This is safe because the split mutual information is symmetric.

• return_all_scores (bool) – If True, returns a dict mapping subset tuples -> MI score (can be large: O(2^K)).

Returns:

result –

• “S”: tuple of indices for the best subset

• ”Sc”: tuple of indices for the complement

• ”mi”: best MI value

• ”cov”: covariance used

• optionally “scores”: dict[(tuple)->float] if return_all_scores=True

Return type:

dict with keys

qpots.utils.tc_utils.cholesky_with_jitter(R, initial_jitter=1e-06, max_jitter=0.01, factor=10.0)

Compute a Cholesky factor, increasing diagonal jitter on failure.

Parameters:

• R (torch.Tensor) – Square covariance or correlation matrix.

• initial_jitter (float, optional) – Initial diagonal jitter to try after the first failure.

• max_jitter (float, optional) – Maximum allowed jitter before raising an error.

• factor (float, optional) – Multiplicative increase applied to the jitter after each failed attempt.

Returns:

Lower-triangular Cholesky factor.

Return type:

torch.Tensor

Raises:

RuntimeError – If the decomposition still fails after the jitter exceeds max_jitter.

qpots.utils.tc_utils.computeTC(x, mt_model)

Compute the total correlation (TC) from the posterior covariance of a multi-output GP model at a given input.

This function evaluates the model posterior at a given point, extracts the covariance matrix of the joint output distribution, and computes the total correlation (a measure of statistical dependence between outputs).

Parameters:

• x (torch.Tensor) – Input tensor. Can be of shape (…, d) or flat. Internally reshaped to -1 x d_eff, where d_eff is one fewer than the model training-input dimension. This assumes the final column of the training inputs is a task index and is not present in x.

• mt_model (MultiTaskGP or compatible model) – A trained multi-output GP model that supports .posterior(X) and returns a multivariate normal distribution with a full covariance matrix across outputs.

Returns:

tc_ – Scalar tensor representing the total correlation of the posterior output distribution at input x.

Return type:

torch.Tensor

Notes

• The covariance matrix reflects dependencies between outputs (e.g., tasks in a MultiTaskGP).

• Total correlation (TC) is a multivariate generalization of mutual information: TC = sum of marginal entropies - joint entropy. It is zero if outputs are independent and positive otherwise.

Assumes corr_and_total_correlation(cov) returns (correlation_matrix, total_correlation).

Assumptions

mt_model.train_inputs[0] must exist, the effective feature dimension is one fewer than the training-input dimension, and the posterior covariance must be small enough to materialize in memory.

qpots.utils.tc_utils.corr_and_total_correlation(cov: Tensor, jitter: float = 1e-06, eps: float = 1e-12)

Compute task correlation matrix R and total correlation TC from an inter-task covariance matrix.

Parameters:

• cov – (…, m, m) symmetric covariance matrix across m tasks (e.g., posterior Cov[f(x)]).

• jitter – diagonal jitter added to R before Cholesky/logdet for numerical stability.

• eps – clamp floor for variances to avoid divide-by-zero.

Returns:

(…, m, m) correlation matrix. TC: (…) total correlation in nats, TC = -0.5 * logdet(R).

Return type:

R

qpots.utils.tc_utils.fit_mtgp(train_X, train_Y, d, n_train, device, dtype)

Fit a two-output MultiTaskGP in long format.

Parameters:

• train_X (torch.Tensor) – Shared input locations with shape n_train x d.

• train_Y (torch.Tensor) – Two-output response matrix with shape n_train x 2.

• d (int) – Number of design variables. The task feature is appended at column d.

• n_train (int) – Number of initial training rows to use.

• device (torch.device or str) – Device on which to build and fit the model.

• dtype (torch.dtype) – Floating point dtype for model tensors.

Returns:

Fitted rank-1 multitask GP model.

Return type:

botorch.models.MultiTaskGP

Notes

This helper is retained for older multitask experiments. New code should prefer ModelObject.fit_multitask_gp because it supports an arbitrary number of objectives/constraints and missing outputs.

qpots.utils.tc_utils.get_model_identified_hv_maximizing_set(model, problem, ref_point, train_y, multiplier=1, max_gen=100, ncons=0)

Construct a Pymoo optimization problem that uses a GP posterior sample to identify a hypervolume-maximizing set via NSGA-II.

This function defines an inner Pymoo Problem class whose objective evaluations are derived from samples of the GP posterior. It is intended to be used with evolutionary algorithms (e.g., NSGA-II) to approximate the Pareto set that maximizes hypervolume under the learned model.

Parameters:

• model (ModelListGP or compatible GP model) – A trained BoTorch model. Must support .posterior(X) and return a posterior over all outputs (objectives + constraints).

• problem (object) –

  A problem definition containing:

  • problem.dim : int, input dimension (d)

  • problem.nobj : int, number of objectives (m)

• ref_point (torch.Tensor) – Reference point for hypervolume computation. Used externally (not directly inside _evaluate), but defines dtype/device and objective dimensionality consistency.

• train_y (torch.Tensor) – Training outputs used to standardize the GP model. Required here to unstandardize posterior samples so that optimization is performed in the original objective space.

• multiplier (int, optional (default=1)) –

  Scales the population size:

  population_size = 100 * dim * multiplier

• max_gen (int, optional (default=100)) – Maximum number of generations for the evolutionary algorithm. (Used externally when running NSGA-II.)

• ncons (int, optional (default=0)) –

  Number of constraint outputs in the model. Assumes constraints are the last ncons outputs.

  Constraint convention: - Feasible if constraint values >= 0 - Infeasible points are penalized in objective space

Returns:

• PosteriorMeanPymooProblem (pymoo.core.problem.Problem) – A configured Pymoo problem instance that can be passed to NSGA-II or other evolutionary algorithms.

• Internal Class (PosteriorMeanPymooProblem)

• —————————————-

• Defines the optimization problem evaluated by NSGA-II.

• Key Features

• ————

• - Decision variables – x ∈ [0, 1]^d (assumes normalized design space)

• - Objective evaluation –

  1. Convert numpy input to torch tensor

  2. Evaluate GP posterior at X

  3. Draw a single sample using Sobol QMC sampler

  4. Unstandardize outputs using train_y

  5. Apply constraint handling (if any)

  6. Return NEGATED objectives (since Pymoo minimizes)

• - Stochastic evaluation – Uses a fixed-seed SobolQMCNormalSampler for reproducibility.

• Constraint Handling

• ——————-

• • Constraints are assumed to be the last ncons outputs.

• - Feasibility condition (all constraints >= 0)

• - Infeasible points are penalized by assigning a large negative – value (-1e12) to their objective values before minimization.

• Important Conventions

• ———————

• • Pymoo performs minimization → objectives are negated (out[“F”] = -f)

• - Model operates in standardized space → outputs are unstandardized

• - Sampling (not mean) is used → enables exploration of model uncertainty

Notes

• Although named “PosteriorMean…”, this actually uses posterior samples, not the mean.

• Hypervolume is not computed directly here; instead, NSGA-II approximates the Pareto front, which can later be evaluated using HV.

• The Sobol sampler ensures deterministic behavior across evaluations when using the same seed.

qpots.utils.tc_utils.mutual_information_split_gaussian(cov: ndarray, S: Sequence[int], *, jitter: float = 1e-10, base: float = 2.718281828459045) → float

Compute I(Y_S ; Y_{Sc}) under a multivariate Gaussian with covariance cov.

Parameters:

• cov ((K, K) array) – Covariance matrix of Y.

• S (sequence of ints) – Indices in the subset S.

• jitter (float) – Diagonal jitter used for stable log-determinants.

• base (float) – Log base. Use base=2.0 for bits, base=np.e for nats.

Returns:

mi – Mutual information I(Y_S ; Y_{Sc}) in the chosen log base.

Return type:

float

qpots.utils.tc_utils.qmaximin(train_X, X, *, q: int = 1, return_index: bool = False, return_distance: bool = False)

Select diverse candidates with greedy sequential maximin sampling.

At each step, this routine picks the candidate whose distance to the closest observed or already-selected point is largest. This is also known as farthest-point sampling and is useful when a sampled Pareto set contains many near-duplicates.

Parameters:

• train_X (array-like or torch.Tensor) – Existing observed design points with shape ... x d.

• X (array-like or torch.Tensor) – Candidate design points with shape ... x d. The final dimension must match train_X.

• q (int, optional) – Number of candidates to select. If fewer candidates are available, all available candidates are returned.

• return_index (bool, optional) – If True, also return index tuples into X.shape[:-1] for the selected candidates.

• return_distance (bool, optional) – If True, also return the maximin distance achieved when each point was selected.

Returns:

Selected points with shape k x d, where k = min(q, num_candidates). Optional indices and distances are appended when requested.

Return type:

torch.Tensor or tuple

qpots.utils.tc_utils.unstandardize_ignore_nan(Y: Tensor, train_y: Tensor, correction: int = 1) → Tensor

Reverse the standardization of output Y using the mean and standard deviation computed from the training data, works when NaN values are in the train_y tensor.

Parameters:

• Y (torch.Tensor) – The standardized output tensor.

• train_y (torch.Tensor) – The training output data used to compute the mean and standard deviation.

Returns:

The unstandardized output tensor.

Return type:

torch.Tensor

qpots.utils.tc_utils.update_mtgp_with_new_data(mt_model: MultiTaskGP, x_new: Tensor, y_new: Tensor, task_feature: int, rank: int = 1, refit_hyperparams: bool = True) → MultiTaskGP

Rebuild a MultiTaskGP after adding partially observed data.

New observations are supplied in wide format with NaN values for outputs that were not evaluated. The function converts those observations to long format, recovers the existing model’s raw training targets, concatenates old and new data, rebuilds a fresh MultiTaskGP, and warm-starts compatible hyperparameters from the previous model.

Parameters:

• mt_model (botorch.models.MultiTaskGP) – Existing fitted multitask model.

• x_new (torch.Tensor) – New design points with shape q x d and no task feature.

• y_new (torch.Tensor) – New outputs with shape q x K. Entries that were not evaluated must be NaN.

• task_feature (int) – Location of the task-feature column in the augmented inputs.

• rank (int, optional) – Rank parameter for the rebuilt MultiTaskGP.

• refit_hyperparams (bool, optional) – If True, refit the marginal log likelihood after rebuilding.

Returns:

Updated multitask model containing old and newly observed data.

Return type:

botorch.models.MultiTaskGP

qpots.utils.tc_utils.wide_to_long_mt(x: Tensor, y: Tensor, task_feature: int) → tuple[Tensor, Tensor]

Convert partially observed wide data into MultiTaskGP long format.

Parameters:

• x (torch.Tensor) – Candidate locations with shape q x d and no task feature.

• y (torch.Tensor) – Output matrix with shape q x K. Missing or intentionally skipped outputs should be represented by NaN.

• task_feature (int) – Column where the task index is inserted in the augmented d + 1 dimensional design matrix.

Returns:

X_long with shape n_obs x (d + 1) and Y_long with shape n_obs x 1. Only finite entries of y are included.

Return type:

tuple[torch.Tensor, torch.Tensor]

Raises:

ValueError – If inputs are not two-dimensional, row counts do not match, the task feature is invalid, or no finite observations are present.