uqtils

Assorted utilities for uncertainty quantification and scientific computing.

• Author - Joshua Eckels (eckelsjd.@umich.edu)
• License - GPL-3.0

Includes:

• MCMC - A standard DRAM MCMC sampler.
• Gradients - Vectorized finite-difference implementation of Jacobian and Hessians.
• Plotting - Some plotting utilities for matplotlib.
• Sobol' - Sobol' global, variance-based sensitivity analysis.

approx_hess(func, theta, pert=0.01)

Approximate Hessian of func at a specified theta location using finite difference approximation.

PARAMETER	DESCRIPTION
func	expects to be called as func(theta) -> (..., y_dim)
theta	(..., theta_dim), points to linearize model about TYPE: Array
pert	perturbation percent for approximate partial derivatives DEFAULT: 0.01

RETURNS	DESCRIPTION
ndarray	(..., y_dim, theta_dim, theta_dim), the approximate Hessian (theta_dim, theta_dim) at all locations (...,) for vector-valued function of dimension y_dim

Source code in src/uqtils/gradient.py

def approx_hess(func, theta: Array, pert=0.01) -> np.ndarray:
    """Approximate Hessian of `func` at a specified `theta` location using finite difference approximation.

    :param func: expects to be called as `func(theta) -> (..., y_dim)`
    :param theta: `(..., theta_dim)`, points to linearize model about
    :param pert: perturbation percent for approximate partial derivatives
    :returns H: `(..., y_dim, theta_dim, theta_dim)`, the approximate Hessian `(theta_dim, theta_dim)` at all locations
                `(...,)` for vector-valued function of dimension `y_dim`
    """
    theta = np.atleast_1d(theta)
    shape = theta.shape[:-1]                # (...,)
    theta_dim = theta.shape[-1]             # Number of parameters
    dtheta = pert * np.abs(theta)

    # Make sure dtheta is not 0 anywhere
    for i in range(theta_dim):
        zero_idx = np.isclose(dtheta[..., i], 0)
        if np.any(zero_idx):
            subs_dtheta = pert * np.abs(np.mean(theta[..., i]))
            if np.isclose(subs_dtheta, 0):
                subs_dtheta = pert
            dtheta[zero_idx, i] = subs_dtheta

    # Return the Hessians (..., y_dim, theta_dim, theta_dim)
    y_dim, H = None, None

    for i in range(theta_dim):
        for j in range(i, theta_dim):
            # Allocate space at 4 grid points (n1=-1, p1=+1)
            theta_n1_n1 = np.copy(theta)
            theta_p1_p1 = np.copy(theta)
            theta_n1_p1 = np.copy(theta)
            theta_p1_n1 = np.copy(theta)

            # Perturbations to theta in each direction
            theta_n1_n1[..., i] -= dtheta[..., i]
            theta_n1_n1[..., j] -= dtheta[..., j]
            f_n1_n1 = func(theta_n1_n1)

            theta_p1_p1[..., i] += dtheta[..., i]
            theta_p1_p1[..., j] += dtheta[..., j]
            f_p1_p1 = func(theta_p1_p1)

            theta_n1_p1[..., i] -= dtheta[..., i]
            theta_n1_p1[..., j] += dtheta[..., j]
            f_n1_p1 = func(theta_n1_p1)

            theta_p1_n1[..., i] += dtheta[..., i]
            theta_p1_n1[..., j] -= dtheta[..., j]
            f_p1_n1 = func(theta_p1_n1)

            if H is None:
                y_dim = f_p1_n1.shape[-1]
                H = np.empty(shape + (y_dim, theta_dim, theta_dim))

            res = (f_n1_n1 + f_p1_p1 - f_n1_p1 - f_p1_n1) / np.expand_dims(4 * dtheta[..., i] * dtheta[..., j],
                                                                           axis=-1)
            H[..., i, j] = res
            H[..., j, i] = res

    if y_dim == 1:
        H = np.squeeze(H, axis=-3)
        if theta_dim == 1:
            H = np.squeeze(H, axis=(-1, -2))
    return np.atleast_1d(H)

approx_jac(func, theta, pert=0.01)

Approximate Jacobian of func at a specified theta location using finite difference approximation.

PARAMETER	DESCRIPTION
func	expects to be called as func(theta) -> (..., y_dim)
theta	(..., theta_dim), points to linearize model about TYPE: Array
pert	perturbation percent for approximate partial derivatives DEFAULT: 0.01

RETURNS	DESCRIPTION
ndarray	(..., y_dim, theta_dim), the approximate Jacobian (y_dim, theta_dim) at all locations (...)

Source code in src/uqtils/gradient.py

def approx_jac(func, theta: Array, pert=0.01) -> np.ndarray:
    """Approximate Jacobian of `func` at a specified `theta` location using finite difference approximation.

    :param func: expects to be called as `func(theta) -> (..., y_dim)`
    :param theta: `(..., theta_dim)`, points to linearize model about
    :param pert: perturbation percent for approximate partial derivatives
    :returns J: `(..., y_dim, theta_dim)`, the approximate Jacobian `(y_dim, theta_dim)` at all locations `(...)`
    """
    theta = np.atleast_1d(theta)
    shape = theta.shape[:-1]                # (...,)
    theta_dim = theta.shape[-1]             # Number of parameters
    dtheta = pert * np.abs(theta)

    # Make sure dtheta is not 0 anywhere
    for i in range(theta_dim):
        zero_idx = np.isclose(dtheta[..., i], 0)
        if np.any(zero_idx):
            subs_dtheta = pert * np.abs(np.mean(theta[..., i]))
            if np.isclose(subs_dtheta, 0):
                subs_dtheta = pert
            dtheta[zero_idx, i] = subs_dtheta

    # Return the Jacobians (..., y_dim, theta_dim)
    J, y_dim = None, None

    for i in range(theta_dim):
        theta_n1 = np.copy(theta)
        theta_p1 = np.copy(theta)

        # Perturbations to theta
        theta_n1[..., i] -= dtheta[..., i]
        theta_p1[..., i] += dtheta[..., i]
        f_n1 = func(theta_n1)
        f_p1 = func(theta_p1)

        if J is None:
            y_dim = f_p1.shape[-1]
            J = np.empty(shape + (y_dim, theta_dim))

        J[..., i] = (f_p1 - f_n1) / np.expand_dims(2 * dtheta[..., i], axis=-1)

    if y_dim == 1:
        J = np.squeeze(J, axis=-2)
        if theta_dim == 1:
            J = np.squeeze(J, axis=-1)
    return np.atleast_1d(J)

autocorrelation(samples, maxlag=100, step=1)

Compute the auto-correlation of a set of samples.

PARAMETER	DESCRIPTION
samples	(niter, nwalk, ndim) samples returned from dram or a similar MCMC routine
maxlag	maximum distance to compute the correlation for DEFAULT: 100
step	step between distances from 0 to maxlag for which to compute the correlations DEFAULT: 1

RETURNS	DESCRIPTION
	lags, autos, tau, ess - the lag times, auto-correlations, integrated auto-correlation, and effective sample sizes

Source code in src/uqtils/mcmc.py

def autocorrelation(samples, maxlag=100, step=1):
    """Compute the auto-correlation of a set of samples.

    :param samples: `(niter, nwalk, ndim)` samples returned from `dram` or a similar MCMC routine
    :param maxlag: maximum distance to compute the correlation for
    :param step: step between distances from 0 to `maxlag` for which to compute the correlations
    :returns: lags, autos, tau, ess - the lag times, auto-correlations, integrated auto-correlation,
              and effective sample sizes
    """
    niter, nwalk, ndim = samples.shape
    mean = np.mean(samples, axis=0)
    var = np.sum((samples - mean[np.newaxis, ...]) ** 2, axis=0)

    lags = np.arange(0, maxlag, step)
    autos = np.zeros((len(lags), nwalk, ndim))
    for zz, lag in enumerate(lags):
        # compute the covariance between all samples *lag apart*
        for ii in range(niter - lag):
            autos[zz, ...] += (samples[ii, ...] - mean) * (samples[ii + lag, ...] - mean)
        autos[zz, ...] /= var
    tau = 1 + 2 * np.sum(autos, axis=0)     # Integrated auto-correlation
    ess = niter / tau                       # Effective sample size
    return lags, autos, tau, ess

ax_default(ax, xlabel='', ylabel='', legend=None, cmap='tab10')

Nice default plt formatting for plotting X-Y data.

PARAMETER	DESCRIPTION
ax	the axes to apply these settings to TYPE: Axes
xlabel	the xlabel to set for ax DEFAULT: ''
ylabel	the ylabel to set for ax DEFAULT: ''
legend	will display a legend if bool(legend) is truthy, can pass a dict of legend kwargs here (optional) DEFAULT: None
cmap	colormap to use for cycling DEFAULT: 'tab10'

Source code in src/uqtils/plots.py

def ax_default(ax: plt.Axes, xlabel='', ylabel='', legend=None, cmap='tab10'):
    """Nice default plt formatting for plotting X-Y data.

    :param ax: the axes to apply these settings to
    :param xlabel: the xlabel to set for `ax`
    :param ylabel: the ylabel to set for `ax`
    :param legend: will display a legend if bool(legend) is truthy, can pass a dict of legend kwargs here (optional)
    :param cmap: colormap to use for cycling
    """
    default_leg = {'fancybox': True, 'facecolor': 'white', 'framealpha': 1, 'loc': 'best', 'edgecolor': 'k'}
    leg_use = legend if isinstance(legend, dict) else default_leg
    for key, val in default_leg.items():
        if key not in leg_use:
            leg_use[key] = val

    ax.set_prop_cycle(_get_cycle(cmap))
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.tick_params(axis='both', which='both', direction='in')
    if legend:
        leg = ax.legend(**leg_use)
        return leg

dram(logpdf, x0, niter, cov0=None, gamma=0.5, eps=1e-06, adapt_after=100, adapt_interval=10, delayed=True, progress=True, filename=None)

Delayed adaptive metropolis-hastings MCMC with a Gaussian proposal.

PARAMETER	DESCRIPTION
logpdf	log PDF function of target distribution
x0	(nwalkers, ndim) initial parameter samples, ignored if samples exist in filename
cov0	(ndim, ndim) the initial proposal covariance, defaults to identity or cov value in filename DEFAULT: None
niter	number of iterations
gamma	scale factor for the covariance matrix for delayed rejection step DEFAULT: 0.5
eps	small constant for making sure covariance is well-conditioned DEFAULT: 1e-06
adapt_after	the number of iterations before covariance adaptation begins (ignored if <=0) DEFAULT: 100
adapt_interval	the number of iterations between each covariance adaptation (ignored if adapt_after<=0) DEFAULT: 10
delayed	whether to try to sample again after first rejection DEFAULT: True
progress	whether to display progress of the sampler DEFAULT: True
filename	if specified, an hdf5 file to save results to. If the file already has dram results, the new samples will be appended. Follows the same format as the emcee library DEFAULT: None

RETURNS	DESCRIPTION
	samples, log_pdf, acceptance - (niter, nwalkers, ndim) samples of the target distribution, the logpdf values at these locations, and the cumulative number of accepted samples per walker

Source code in src/uqtils/mcmc.py

def dram(logpdf, x0, niter, cov0=None, gamma=0.5, eps=1e-6, adapt_after=100, adapt_interval=10,
         delayed=True, progress=True, filename=None):
    """Delayed adaptive metropolis-hastings MCMC with a Gaussian proposal.

    :param logpdf: log PDF function of target distribution
    :param x0: `(nwalkers, ndim)` initial parameter samples, ignored if samples exist in `filename`
    :param cov0: `(ndim, ndim)` the initial proposal covariance, defaults to identity or `cov` value in filename
    :param niter: number of iterations
    :param gamma: scale factor for the covariance matrix for delayed rejection step
    :param eps: small constant for making sure covariance is well-conditioned
    :param adapt_after: the number of iterations before covariance adaptation begins (ignored if <=0)
    :param adapt_interval: the number of iterations between each covariance adaptation (ignored if `adapt_after<=0`)
    :param delayed: whether to try to sample again after first rejection
    :param progress: whether to display progress of the sampler
    :param filename: if specified, an hdf5 file to save results to. If the file already has dram results, the new
                     samples will be appended. Follows the same format as the `emcee` library
    :returns: `samples, log_pdf, acceptance` - `(niter, nwalkers, ndim)` samples of the target distribution, the logpdf
              values at these locations, and the cumulative number of accepted samples per walker
    """
    # Override x0, cov0 if filename already has samples
    try:
        if filename is not None:
            with h5py.File(filename, 'a') as fd:
                group = fd.get('mcmc', None)
                if group is not None:
                    x0 = group['chain'][-1, ...]
                    if cov0 is None:
                        cov0 = np.array(group['cov'])  # only override if cov0 is not passed in
                    niter += 1
    except Exception as e:
        warnings.warn(str(e))

    # Initialize
    x0 = np.atleast_2d(x0)
    nwalk, ndim = x0.shape
    cov0 = np.eye(ndim) if cov0 is None else cov0
    sd = (2.4**2/ndim)
    curr_cov = np.broadcast_to(cov0, (nwalk, ndim, ndim)).copy().astype(x0.dtype)
    curr_chol = np.linalg.cholesky(curr_cov)
    adapt_cov = curr_cov.copy()  # adaptive covariance
    curr_mean = x0
    curr_loc_logpdf = logpdf(x0)
    samples = np.empty((niter, nwalk, ndim), dtype=x0.dtype)
    log_pdf = np.empty((niter, nwalk), dtype=x0.dtype)
    accepted = np.zeros((nwalk,), dtype=x0.dtype)
    samples[0, ...] = x0
    log_pdf[0, ...] = curr_loc_logpdf

    def accept_first(curr_log, prop_log):
        with np.errstate(over='ignore'):
            # Overflow values go to -> infty, so they will always get accepted
            ret = np.minimum(1.0, np.exp(prop_log - curr_log))
        return ret

    # Main sample loop
    iterable = tqdm.tqdm(range(niter-1)) if progress else range(niter-1)
    # --8<-- [start:dram]
    for i in iterable:
        # Propose sample
        x1 = samples[i, ...]
        y1 = normal_sample(x1, curr_chol, sqrt=True)    # (nwalkers, ndim)
        x1_log = curr_loc_logpdf
        y1_log = logpdf(y1)

        # Compute first acceptance
        with np.errstate(invalid='ignore'):
            a1 = y1_log - x1_log                        # (nwalkers,)
        a1_idx = a1 > 0
        a1_idx |= np.log(np.random.rand(nwalk)) < a1
        samples[i + 1, a1_idx, :] = y1[a1_idx, :]
        samples[i + 1, ~a1_idx, :] = x1[~a1_idx, :]
        curr_loc_logpdf[a1_idx] = y1_log[a1_idx]
        accepted[a1_idx] += 1

        # Second level proposal
        if delayed and np.any(~a1_idx):
            y2 = normal_sample(x1[~a1_idx, :], curr_chol[~a1_idx, ...] * np.sqrt(gamma), sqrt=True)
            y2_log = logpdf(y2)
            with ((np.errstate(divide='ignore', invalid='ignore'))):
                # If a(y2, y1)=1, then log(1-a(y2,y1)) -> -infty and a2 -> 0
                frac_1 = y2_log - x1_log[~a1_idx]
                frac_2 = (normal_pdf(y1[~a1_idx, :], y2, curr_cov[~a1_idx, ...], logpdf=True) -
                          normal_pdf(y1[~a1_idx, :], x1[~a1_idx, :], curr_cov[~a1_idx, ...], logpdf=True))
                frac_3 = (np.log(1 - accept_first(y2_log, y1_log[~a1_idx])) -
                          np.log(1 - np.minimum(1.0, np.exp(a1[~a1_idx]))))
                a2 = frac_1 + frac_2 + frac_3
            a2_idx = a2 > 0
            a2_idx |= np.log(np.random.rand(a2.shape[0])) < a2

            sample_a2_idx = np.where(~a1_idx)[0][a2_idx]  # Indices that were False the 1st time, then true the 2nd
            samples[i + 1, sample_a2_idx, :] = y2[a2_idx, :]
            curr_loc_logpdf[sample_a2_idx] = y2_log[a2_idx]
            accepted[sample_a2_idx] += 1

        log_pdf[i+1, ...] = curr_loc_logpdf

        # Update the sample mean and cov every iteration
        if adapt_after > 0:
            k = i + 1
            last_mean = curr_mean.copy()
            curr_mean = (1/(k+1)) * samples[k, ...] + (k/(k+1))*last_mean
            mult = (np.eye(ndim) * eps + k * last_mean[..., np.newaxis] @ last_mean[..., np.newaxis, :] -
                    (k + 1) * curr_mean[..., np.newaxis] @ curr_mean[..., np.newaxis, :] +
                    samples[k, ..., np.newaxis] @ samples[k, ..., np.newaxis, :])
            adapt_cov = ((k - 1) / k) * adapt_cov + (sd / k) * mult

            if k > adapt_after and k % adapt_interval == 0:
                try:
                    curr_chol[:] = np.linalg.cholesky(adapt_cov)
                    curr_cov[:] = adapt_cov[:]
                except np.linalg.LinAlgError:
                    warnings.warn(f"Non-PSD matrix at k={k}. Ignoring...")
    # --8<-- [end:dram]

    try:
        if filename is not None:
            with h5py.File(filename, 'a') as fd:
                group = fd.get('mcmc', None)
                if group is not None:
                    samples = np.concatenate((group['chain'], samples[1:, ...]), axis=0)
                    log_pdf = np.concatenate((group['log_pdf'], log_pdf[1:, ...]), axis=0)
                    accepted += group['accepted']
                    del group['chain']
                    del group['log_pdf']
                    del group['accepted']
                    del group['cov']
                fd.create_dataset('mcmc/chain', data=samples)
                fd.create_dataset('mcmc/log_pdf', data=log_pdf)
                fd.create_dataset('mcmc/accepted', data=accepted)
                fd.create_dataset('mcmc/cov', data=curr_cov)
    except Exception as e:
        warnings.warn(str(e))

    return samples, log_pdf, accepted

format_input(x, ndim)

Helper function to make sure input x is an ndarray of shape (..., ndim).

PARAMETER	DESCRIPTION
x	if 1d-like as (n,), then converted to 2d as (1, n) if n==ndim or (n, 1) if ndim==1 TYPE: Array
ndim	the dimension of the inputs TYPE: int

RETURNS	DESCRIPTION
tuple[bool, ndarray]	x as at least a 2d array (..., ndim), and whether x was originally 1d-like

Source code in src/uqtils/uq_types.py

def format_input(x: Array, ndim: int) -> tuple[bool, np.ndarray]:
    """Helper function to make sure input `x` is an `ndarray` of shape `(..., ndim)`.

    :param x: if 1d-like as `(n,)`, then converted to 2d as `(1, n) if n==ndim or (n, 1) if ndim==1`
    :param ndim: the dimension of the inputs
    :returns: `x` as at least a 2d array `(..., ndim)`, and whether `x` was originally 1d-like
    """
    x = np.atleast_1d(x)
    is_1d = len(x.shape) == 1
    if is_1d:
        if x.shape[0] != ndim and ndim > 1:
            raise ValueError(f'Input x shape {x.shape} is incompatible with ndim of {ndim}')
        x = np.expand_dims(x, axis=0 if x.shape[0] == ndim else 1)

    return is_1d, x

is_positive_definite(A)

Returns true when input is positive-definite, via Cholesky.

Source code in src/uqtils/mcmc.py

def is_positive_definite(A):
    """Returns true when input is positive-definite, via Cholesky."""
    try:
        _ = np.linalg.cholesky(A)
        return True
    except np.linalg.LinAlgError:
        return False

ishigami(x, a=7.0, b=0.1)

For testing Sobol indices: Ishigami function

Source code in src/uqtils/sobol.py

def ishigami(x, a=7.0, b=0.1):
    """For testing Sobol indices: [Ishigami function](https://doi.org/10.1109/ISUMA.1990.151285)"""
    return {'y': np.sin(x[..., 0:1]) + a*np.sin(x[..., 1:2])**2 + b*(x[..., 2:3]**4)*np.sin(x[..., 0:1])}

ndscatter(samples, labels=None, tick_fmts=None, plot1d=None, plot2d='scatter', cmap='viridis', bins=20, cmin=0, z=None, cb_label=None, cb_norm='linear', subplot_size=3, cov_overlay=None)

Triangle scatter plots of n-dimensional samples.

Warning

Best for dim < 10. You can shrink the subplot_size to assist graphics loading time.

PARAMETER	DESCRIPTION
samples	(N, dim) samples to plot TYPE: ndarray
labels	list of axis labels of length dim TYPE: list[str] DEFAULT: None
tick_fmts	list of str.format() specifiers for ticks, e.g ['{x: ^10.2f}', ...], of length dim TYPE: list[str] DEFAULT: None
plot1d	'hist' or 'kde' for 1d marginals, defaults to plot2d if None TYPE: Literal['kde', 'hist'] DEFAULT: None
plot2d	'hist' for 2d hist plot, 'kde' for kernel density estimation, 'hex', or 'scatter' (default) TYPE: Literal['scatter', 'kde', 'hist', 'hex'] DEFAULT: 'scatter'
cmap	the matplotlib string specifier of a colormap DEFAULT: 'viridis'
bins	number of bins in each dimension for histogram marginals DEFAULT: 20
cmin	the minimum bin count below which the bins are not displayed DEFAULT: 0
z	(N,) a performance metric corresponding to samples, used to color code the scatter plot if provided TYPE: ndarray DEFAULT: None
cb_label	label for color bar (if z is provided) DEFAULT: None
cb_norm	str or plt.colors.Normalize, normalization method for plotting z on scatter plot DEFAULT: 'linear'
subplot_size	size in inches of a single 2d marginal subplot DEFAULT: 3
cov_overlay	(ndim, ndim) a covariance matrix to overlay as a Gaussian kde over the samples DEFAULT: None

RETURNS	DESCRIPTION
	the plt Figure and Axes objects, (returns an additional cb_fig, cb_ax if z is specified)

Source code in src/uqtils/plots.py

def ndscatter(samples: np.ndarray, labels: list[str] = None, tick_fmts: list[str] = None,
              plot1d: Literal['kde', 'hist'] = None, plot2d: Literal['scatter', 'kde', 'hist', 'hex'] = 'scatter',
              cmap='viridis', bins=20, cmin=0, z: np.ndarray = None, cb_label=None, cb_norm='linear',
              subplot_size=3, cov_overlay=None):
    """Triangle scatter plots of n-dimensional samples.

    !!! Warning
        Best for `dim < 10`. You can shrink the `subplot_size` to assist graphics loading time.

    :param samples: `(N, dim)` samples to plot
    :param labels: list of axis labels of length `dim`
    :param tick_fmts: list of str.format() specifiers for ticks, e.g `['{x: ^10.2f}', ...]`, of length `dim`
    :param plot1d: 'hist' or 'kde' for 1d marginals, defaults to plot2d if None
    :param plot2d: 'hist' for 2d hist plot, 'kde' for kernel density estimation, 'hex', or 'scatter' (default)
    :param cmap: the matplotlib string specifier of a colormap
    :param bins: number of bins in each dimension for histogram marginals
    :param cmin: the minimum bin count below which the bins are not displayed
    :param z: `(N,)` a performance metric corresponding to `samples`, used to color code the scatter plot if provided
    :param cb_label: label for color bar (if `z` is provided)
    :param cb_norm: `str` or `plt.colors.Normalize`, normalization method for plotting `z` on scatter plot
    :param subplot_size: size in inches of a single 2d marginal subplot
    :param cov_overlay: `(ndim, ndim)` a covariance matrix to overlay as a Gaussian kde over the samples
    :returns fig, axs: the `plt` Figure and Axes objects, (returns an additional `cb_fig, cb_ax` if `z` is specified)
    """
    N, dim = samples.shape
    x_min = np.min(samples, axis=0)
    x_max = np.max(samples, axis=0)
    show_colorbar = z is not None
    if labels is None:
        labels = [f"x{i}" for i in range(dim)]
    if z is None:
        z = plt.get_cmap(cmap)([0])
    if cb_label is None:
        cb_label = 'Performance metric'

    def tick_format_func(value, pos):
        if np.isclose(value, 0):
            return f'{value:.2f}'
        if abs(value) > 1000:
            return f'{value:.2E}'
        if abs(value) > 100:
            return f'{int(value):d}'
        if abs(value) > 1:
            return f'{value:.2f}'
        if abs(value) > 0.01:
            return f'{value:.4f}'
        if abs(value) < 0.01:
            return f'{value:.2E}'
    default_ticks = FuncFormatter(tick_format_func)
    # if tick_fmts is None:
    #     tick_fmts = ['{x:.2G}' for i in range(dim)]

    # Set up triangle plot formatting
    fig, axs = plt.subplots(dim, dim, sharex='col', sharey='row')
    for i in range(dim):
        for j in range(dim):
            ax = axs[i, j]
            if i == j:                      # 1d marginals on diagonal
                # ax.get_shared_y_axes().remove(ax)
                ax._shared_axes['y'].remove(ax)
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.spines['left'].set_visible(False)
                if i == 0:
                    ax.get_yaxis().set_ticks([])
            if j > i:                       # Clear the upper triangle
                ax.axis('off')
            if i == dim - 1:                # Bottom row
                ax.set_xlabel(labels[j])
                ax.xaxis.set_major_locator(AutoLocator())
                formatter = StrMethodFormatter(tick_fmts[j]) if tick_fmts is not None else default_ticks
                ax.xaxis.set_major_formatter(formatter)
            if j == 0 and i > 0:            # Left column
                ax.set_ylabel(labels[i])
                ax.yaxis.set_major_locator(AutoLocator())
                formatter = StrMethodFormatter(tick_fmts[i]) if tick_fmts is not None else default_ticks
                ax.yaxis.set_major_formatter(formatter)

    if cov_overlay is not None:
        x_overlay = normal_sample(np.mean(samples, axis=0), cov_overlay, 5000)

    # Plot marginals
    for i in range(dim):
        for j in range(dim):
            ax = axs[i, j]
            if i == j:                      # 1d marginals (on diagonal)
                c = plt.get_cmap(cmap)(0)
                plot = plot1d if plot1d is not None else plot2d
                if plot == 'kde':
                    kernel = st.gaussian_kde(samples[:, i])
                    x = np.linspace(x_min[i], x_max[i], 500)
                    ax.fill_between(x, y1=kernel(x), y2=0, lw=0, alpha=0.3, facecolor=c)
                    ax.plot(x, kernel(x), ls='-', c=c, lw=1.5)
                else:
                    ax.hist(samples[:, i], edgecolor='black', color=c, density=True, alpha=0.5,
                            linewidth=1.2, bins=bins)
                if cov_overlay is not None:
                    kernel = st.gaussian_kde(x_overlay[:, i])
                    x = np.linspace(x_min[i], x_max[i], 500)
                    ax.fill_between(x, y1=kernel(x), y2=0, lw=0, alpha=0.5, facecolor=[0.5, 0.5, 0.5])
                    ax.plot(x, kernel(x), ls='-', c='k', lw=1.5, alpha=0.5)
                bottom, top = ax.get_ylim()
                ax.set_ylim([0, top])
            if j < i:                       # 2d marginals (lower triangle)
                ax.set_xlim([x_min[j], x_max[j]])
                ax.set_ylim([x_min[i], x_max[i]])
                if plot2d == 'scatter':
                    sc = ax.scatter(samples[:, j], samples[:, i], s=1.5, c=z, cmap=cmap, norm=cb_norm)
                elif plot2d == 'hist':
                    ax.hist2d(samples[:, j], samples[:, i], bins=bins, cmap=cmap, cmin=cmin)
                elif plot2d == 'kde':
                    kernel = st.gaussian_kde(samples[:, [j, i]].T)
                    xg, yg = np.meshgrid(np.linspace(x_min[j], x_max[j], 40), np.linspace(x_min[i], x_max[i], 40))
                    x = np.vstack([xg.ravel(), yg.ravel()])
                    zg = np.reshape(kernel(x), xg.shape)
                    cs = ax.contourf(xg, yg, zg, 5, cmap=cmap, alpha=0.9, extend='both')
                    cs.cmap.set_under('white')
                    cs.changed()
                    ax.contour(xg, yg, zg, 5, colors=[(0.5, 0.5, 0.5)], linewidths=1.2)
                elif plot2d == 'hex':
                    ax.hexbin(samples[:, j], samples[:, i], gridsize=bins, cmap=cmap, mincnt=cmin)
                else:
                    raise NotImplementedError('This plot type is not known. plot2d=["hist", "kde", "scatter"]')

                if cov_overlay is not None:
                    kernel = st.gaussian_kde(x_overlay[:, [j, i]].T)
                    xg, yg = np.meshgrid(np.linspace(x_min[j], x_max[j], 40), np.linspace(x_min[i], x_max[i], 40))
                    x = np.vstack([xg.ravel(), yg.ravel()])
                    zg = np.reshape(kernel(x), xg.shape)
                    ax.contourf(xg, yg, zg, 4, cmap='Greys', alpha=0.4)
                    ax.contour(xg, yg, zg, 4, colors='k', linewidths=1.5, alpha=0.6)

    fig.set_size_inches(subplot_size * dim, subplot_size * dim)
    fig.tight_layout()

    # Plot colorbar in standalone figure
    if show_colorbar and plot2d == 'scatter':
        cb_fig, cb_ax = plt.subplots(figsize=(1.5, 6))
        cb_fig.subplots_adjust(right=0.7)
        cb_fig.colorbar(sc, cax=cb_ax, orientation='vertical', label=cb_label)
        cb_fig.tight_layout()
        return fig, axs, cb_fig, cb_ax

    return fig, axs

nearest_positive_definite(A)

Find the nearest positive-definite matrix to input.

A Python port of John D'Errico's nearestSPD MATLAB code [1], which credits [2]. [1] https://www.mathworks.com/matlabcentral/fileexchange/42885-nearestspd [2] N.J. Higham, "Computing a nearest symmetric positive semidefinite matrix", 1988

Source code in src/uqtils/mcmc.py

def nearest_positive_definite(A):
    """Find the nearest positive-definite matrix to input.

    A Python port of John D'Errico's `nearestSPD` MATLAB code [1], which credits [2].
    [1] https://www.mathworks.com/matlabcentral/fileexchange/42885-nearestspd
    [2] N.J. Higham, "Computing a nearest symmetric positive semidefinite matrix", 1988
    """
    B = (A + A.T) / 2
    _, s, V = np.linalg.svd(B)
    H = np.dot(V.T, np.dot(np.diag(s), V))
    A2 = (B + H) / 2
    A3 = (A2 + A2.T) / 2
    if is_positive_definite(A3):
        return A3
    spacing = np.spacing(np.linalg.norm(A))
    # The above is different from [1]. It appears that MATLAB's `chol` Cholesky
    # decomposition will accept matrixes with exactly 0-eigenvalue, whereas
    # Numpy's will not. So where [1] uses `eps(mineig)` (where `eps` is Matlab
    # for `np.spacing`), we use the above definition. CAVEAT: our `spacing`
    # will be much larger than [1]'s `eps(mineig)`, since `mineig` is usually on
    # the order of 1e-16, and `eps(1e-16)` is on the order of 1e-34, whereas
    # `spacing` will, for Gaussian random matrixes of small dimension, be on
    # othe order of 1e-16. In practice, both ways converge, as the unit test
    # below suggests.
    eye = np.eye(A.shape[0])
    k = 1
    while not is_positive_definite(A3):
        mineig = np.min(np.real(np.linalg.eigvals(A3)))
        A3 += eye * (-mineig * k ** 2 + spacing)
        k += 1

    return A3

normal_pdf(x, mean, cov, logpdf=False)

Compute the Gaussian pdf at each x location (pretty much however you want).

PARAMETER	DESCRIPTION
x	(..., dim), the locations to evaluate the pdf at TYPE: Array
mean	(..., dim), expected values, where dim is the random variable dimension TYPE: Array
cov	(..., dim, dim), covariance matrices TYPE: Array
logpdf	whether to return the logpdf instead TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
ndarray	(...,) the pdf values at x

Source code in src/uqtils/mcmc.py

def normal_pdf(x: Array, mean: Array, cov: Array, logpdf: bool = False) -> np.ndarray:
    """Compute the Gaussian pdf at each `x` location (pretty much however you want).

    :param x: `(..., dim)`, the locations to evaluate the pdf at
    :param mean: `(..., dim)`, expected values, where dim is the random variable dimension
    :param cov: `(..., dim, dim)`, covariance matrices
    :param logpdf: whether to return the logpdf instead
    :returns: `(...,)` the pdf values at `x`
    """
    x = np.atleast_1d(x)
    mean = np.atleast_1d(mean)
    cov = np.atleast_2d(cov)
    dim = cov.shape[-1]

    preexp = 1 / ((2*np.pi)**(dim/2) * np.linalg.det(cov)**(1/2))
    diff = x - mean
    diff_col = np.expand_dims(diff, axis=-1)    # (..., dim, 1)
    diff_row = np.expand_dims(diff, axis=-2)    # (..., 1, dim)
    inexp = np.squeeze(diff_row @ np.linalg.inv(cov) @ diff_col, axis=(-1, -2))

    pdf = np.log(preexp) - 0.5 * inexp if logpdf else preexp * np.exp(-0.5 * inexp)

    return pdf

normal_sample(mean, cov, size=(), sqrt=False)

Generic batch sample multivariate normal distributions (pretty much however you want).

Note

The provided mean and cov should match along the last dimension, that is the dimension of the random variables to sample. If you want to sample a 1d Gaussian, then you can specify both the mean and covariance as scalars. However, as long as the mean and covariance are broadcastable in size, then you can use this function however you want, (i.e. sample many multivariate distributions at once, all with different means and covariances, etc., just get creative)

PARAMETER	DESCRIPTION
mean	(..., dim), expected values, where dim is the random variable dimension TYPE: Array
cov	(..., dim, dim), covariance matrices (or the sqrt(cov) if sqrt=True) TYPE: Array
size	shape of additional samples TYPE: tuple | int DEFAULT: ()
sqrt	whether cov was passed in already as the sqrt(cov) via cholesky decomposition DEFAULT: False

RETURNS	DESCRIPTION
ndarray	(*size, ..., dim), samples from multivariate distributions

Source code in src/uqtils/mcmc.py

def normal_sample(mean: Array, cov: Array, size: tuple | int = (), sqrt=False) -> np.ndarray:
    """Generic batch sample multivariate normal distributions (pretty much however you want).

    !!! Note
        The provided `mean` and `cov` should match along the last dimension, that is the dimension of the random
        variables to sample. If you want to sample a 1d Gaussian, then you can specify both the mean and covariance
        as scalars. However, as long as the mean and covariance are broadcastable in size, then you can use this
        function however you want, (i.e. sample many multivariate distributions at once, all with different means
        and covariances, etc., just get creative)

    :param mean: `(..., dim)`, expected values, where `dim` is the random variable dimension
    :param cov: `(..., dim, dim)`, covariance matrices (or the sqrt(cov) if `sqrt=True`)
    :param size: shape of additional samples
    :param sqrt: whether `cov` was passed in already as the `sqrt(cov)` via cholesky decomposition
    :returns samples: `(*size, ..., dim)`, samples from multivariate distributions
    """
    mean = np.atleast_1d(mean)
    cov = np.atleast_2d(cov)
    sqrt_cov = cov if sqrt else np.linalg.cholesky(cov)

    if isinstance(size, int):
        size = (size, )
    shape = size + np.broadcast_shapes(mean.shape, cov.shape[:-1])
    x_normal = np.random.standard_normal((*shape, 1)).astype(mean.dtype)
    samples = np.squeeze(sqrt_cov @ x_normal, axis=-1) + mean
    return samples

plot_slice(funs, bds, x0=None, x_idx=None, y_idx=None, N=50, random_walk=False, xlabels=None, ylabels=None, cmap='viridis', fun_labels=None)

Helper function to plot 1d slices of a function(s) over inputs.

PARAMETER	DESCRIPTION
funs	function callable as y=f(x), with x as (..., xdim) and y as (..., ydim), can also be a list of functions to evaluate and plot together.
bds	list of tuples of (min, max) specifying the bounds of the inputs TYPE: list[tuple]
x0	the default values for all inputs; defaults to middle of bds TYPE: Array DEFAULT: None
x_idx	list of input indices to take 1d slices of TYPE: list[int] DEFAULT: None
y_idx	list of output indices to plot 1d slices of TYPE: list[int] DEFAULT: None
N	the number of points to take in each 1d slice TYPE: int DEFAULT: 50
random_walk	whether to slice in a random d-dimensional direction or hold all params const while slicing TYPE: bool DEFAULT: False
xlabels	list of labels for the inputs TYPE: list[str] DEFAULT: None
ylabels	list of labels for the outputs TYPE: list[str] DEFAULT: None
cmap	the name of the matplotlib colormap to use DEFAULT: 'viridis'
fun_labels	the legend labels if plotting multiple functions on each plot DEFAULT: None

RETURNS	DESCRIPTION
	fig, ax with num_inputs by num_outputs subplots

Source code in src/uqtils/plots.py

def plot_slice(funs, bds: list[tuple], x0: Array = None, x_idx: list[int] = None,
               y_idx: list[int] = None, N: int = 50, random_walk: bool = False, xlabels: list[str] = None,
               ylabels: list[str] = None, cmap='viridis', fun_labels=None):
    """Helper function to plot 1d slices of a function(s) over inputs.

    :param funs: function callable as `y=f(x)`, with `x` as `(..., xdim)` and `y` as `(..., ydim)`, can also be a list
                of functions to evaluate and plot together.
    :param bds: list of tuples of `(min, max)` specifying the bounds of the inputs
    :param x0: the default values for all inputs; defaults to middle of `bds`
    :param x_idx: list of input indices to take 1d slices of
    :param y_idx: list of output indices to plot 1d slices of
    :param N: the number of points to take in each 1d slice
    :param random_walk: whether to slice in a random d-dimensional direction or hold all params const while slicing
    :param xlabels: list of labels for the inputs
    :param ylabels: list of labels for the outputs
    :param cmap: the name of the matplotlib colormap to use
    :param fun_labels: the legend labels if plotting multiple functions on each plot
    :returns: `fig, ax` with `num_inputs` by `num_outputs` subplots
    """
    funs = funs if isinstance(funs, list) else [funs]
    x_idx = list(np.arange(0, min(3, len(bds)))) if x_idx is None else x_idx
    y_idx = [0] if y_idx is None else y_idx
    xlabels = [f'x{i}' for i in range(len(x_idx))] if xlabels is None else xlabels
    ylabels = [f'QoI {i}' for i in range(len(y_idx))] if ylabels is None else ylabels
    fun_labels = [f'fun {i}' for i in range(len(funs))] if fun_labels is None else fun_labels
    x0 = [(b[0] + b[1]) / 2 for b in bds] if x0 is None else x0
    x0 = np.atleast_1d(x0)
    xdim = x0.shape[0]
    lb = np.atleast_1d([b[0] for b in bds])
    ub = np.atleast_1d([b[1] for b in bds])
    cmap = plt.get_cmap(cmap)

    # Construct sliced inputs
    xs = np.zeros((N, len(x_idx), xdim))
    for i in range(len(x_idx)):
        if random_walk:
            # Make a random straight-line walk across d-cube
            r0 = np.random.rand(xdim) * (ub - lb) + lb
            r0[x_idx[i]] = lb[x_idx[i]]                     # Start slice at this lower bound
            rf = np.random.rand(xdim) * (ub - lb) + lb
            rf[x_idx[i]] = ub[x_idx[i]]                     # Slice up to this upper bound
            xs[0, i, :] = r0
            for k in range(1, N):
                xs[k, i, :] = xs[k-1, i, :] + (rf-r0)/(N-1)
        else:
            # Otherwise, only slice one variable
            for j in range(xdim):
                if j == x_idx[i]:
                    xs[:, i, j] = np.linspace(lb[x_idx[i]], ub[x_idx[i]], N)
                else:
                    xs[:, i, j] = x0[j]

    # Compute function values and show ydim by xdim grid of subplots
    ys = []
    for func in funs:
        y = func(xs)
        if y.shape == (N, len(x_idx)):
            y = y[..., np.newaxis]
        ys.append(y)
    c_intervals = np.linspace(0, 1, len(ys))

    fig, axs = plt.subplots(len(y_idx), len(x_idx), sharex='col', sharey='row')
    for i in range(len(y_idx)):
        for j in range(len(x_idx)):
            if len(y_idx) == 1:
                ax = axs if len(x_idx) == 1 else axs[j]
            elif len(x_idx) == 1:
                ax = axs if len(y_idx) == 1 else axs[i]
            else:
                ax = axs[i, j]
            x = xs[:, j, x_idx[j]]
            for k in range(len(ys)):
                y = ys[k][:, j, y_idx[i]]
                ax.plot(x, y, ls='-', color=cmap(c_intervals[k]), label=fun_labels[k])
            ylabel = ylabels[i] if j == 0 else ''
            xlabel = xlabels[j] if i == len(y_idx) - 1 else ''
            legend = (i == 0 and j == len(x_idx) - 1 and len(ys) > 1)
            ax_default(ax, xlabel, ylabel, legend=legend)
    fig.set_size_inches(3 * len(x_idx), 3 * len(y_idx))
    fig.tight_layout()

    return fig, axs

sobol_sa(model, sampler, num_samples, qoi_idx=None, qoi_labels=None, param_labels=None, plot=False, verbose=True, cmap='viridis', compute_s2=False)

Perform a global Sobol' sensitivity analysis.

PARAMETER	DESCRIPTION
model	callable as y=model(x), with y=(..., ydim), x=(..., xdim)
sampler	callable as x=sampler(shape), with x=(*shape, xdim)
num_samples	number of samples TYPE: int
qoi_idx	list of indices of model output to report results for TYPE: list[int] DEFAULT: None
qoi_labels	list of labels for plotting QoIs TYPE: list[str] DEFAULT: None
param_labels	list of labels for plotting input parameters TYPE: list[str] DEFAULT: None
plot	whether to plot bar/pie charts TYPE: bool DEFAULT: False
verbose	whether to print S1/ST/S2 results to the console TYPE: bool DEFAULT: True
cmap	str specifier of plt.colormap for bar/pie charts TYPE: str DEFAULT: 'viridis'
compute_s2	whether to compute the second order indices TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
	S1, [S2], ST, the first, second, and total order Sobol' indices

Source code in src/uqtils/sobol.py

def sobol_sa(model, sampler, num_samples: int, qoi_idx: list[int] = None, qoi_labels: list[str] = None,
             param_labels: list[str] = None, plot: bool = False, verbose: bool = True, cmap: str = 'viridis',
             compute_s2: bool = False):
    """Perform a global Sobol' sensitivity analysis.

    :param model: callable as `y=model(x)`, with `y=(..., ydim)`, `x=(..., xdim)`
    :param sampler: callable as `x=sampler(shape)`, with `x=(*shape, xdim)`
    :param num_samples: number of samples
    :param qoi_idx: list of indices of model output to report results for
    :param qoi_labels: list of labels for plotting QoIs
    :param param_labels: list of labels for plotting input parameters
    :param plot: whether to plot bar/pie charts
    :param verbose: whether to print `S1/ST/S2` results to the console
    :param cmap: `str` specifier of `plt.colormap` for bar/pie charts
    :param compute_s2: whether to compute the second order indices
    :return: `S1`, `[S2]`, `ST`, the first, second, and total order Sobol' indices
    """
    # Get sample matrices (N, xdim)
    A = sampler((num_samples,))
    B = sampler((num_samples,))
    xdim = A.shape[-1]
    AB = np.tile(np.expand_dims(A, axis=-2), (1, xdim, 1))
    BA = np.tile(np.expand_dims(B, axis=-2), (1, xdim, 1))
    for i in range(xdim):
        AB[:, i, i] = B[:, i]
        BA[:, i, i] = A[:, i]

    # Evaluate the model; (xdim+2)*N evaluations required
    fA = model(A)       # (N, ydim)
    fB = model(B)       # (N, ydim)
    fAB = model(AB)     # (N, xdim, ydim)
    fBA = model(BA)     # (N, xdim, ydim)
    ydim = fA.shape[-1]

    # Normalize model outputs to N(0, 1) for better stability
    Y = np.concatenate((fA, fB, fAB.reshape((-1, ydim)), fBA.reshape((-1, ydim))), axis=0)
    mu, std = np.mean(Y, axis=0), np.std(Y, axis=0)
    fA = (fA - mu) / std
    fB = (fB - mu) / std
    fAB = (fAB - mu) / std
    fBA = (fBA - mu) / std

    # Compute sensitivity indices
    vY = np.var(np.concatenate((fA, fB), axis=0), axis=0)   # (ydim,)
    fA = np.expand_dims(fA, axis=-2)                        # (N, 1, ydim)
    fB = np.expand_dims(fB, axis=-2)                        # (N, 1, ydim)
    S1 = fB * (fAB - fA) / vY                               # (N, xdim, ydim)
    ST = 0.5 * (fA - fAB)**2 / vY                           # (N, xdim, ydim)

    # Second-order indices
    if compute_s2:
        Vij = np.expand_dims(fBA, axis=2) * np.expand_dims(fAB, axis=1) - \
              np.expand_dims(fA, axis=1) * np.expand_dims(fB, axis=1)   # (N, xdim, xdim, ydim)
        si = fB * (fAB - fA)
        Vi = np.expand_dims(si, axis=2)
        Vj = np.expand_dims(si, axis=1)
        S2 = (Vij - Vi - Vj) / vY                                       # (N, xdim, xdim, ydim)
        S2_est = np.mean(S2, axis=0)
        S2_se = np.sqrt(np.var(S2, axis=0) / num_samples)

    # Get mean values and MC error
    S1_est = np.mean(S1, axis=0)
    S1_se = np.sqrt(np.var(S1, axis=0) / num_samples)
    ST_est = np.mean(ST, axis=0)
    ST_se = np.sqrt(np.var(ST, axis=0) / num_samples)

    # Set defaults for qoi indices/labels
    if qoi_idx is None:
        qoi_idx = list(np.arange(ydim))
    if qoi_labels is None:
        qoi_labels = [f'QoI {i}' for i in range(len(qoi_idx))]
    if param_labels is None:
        param_labels = [f'x{i}' for i in range(xdim)]

    # Print results
    if verbose:
        print(f'{"QoI":>10} {"Param":>10} {"S1_mean":>10} {"S1_err":>10} {"ST_mean":>10} {"ST_err":>10}')
        for i in range(len(qoi_idx)):
            for j in range(xdim):
                q = qoi_idx[i]
                print(f'{qoi_labels[i]:>10} {param_labels[j]:>10} {S1_est[j, q]: 10.3f} {S1_se[j, q]: 10.3f} '
                      f'{ST_est[j, q]: 10.3f} {ST_se[j, q]: 10.3f}')

        if compute_s2:
            print(f'\n{"QoI":>10} {"2nd-order":>20} {"S2_mean":>10} {"S2_err":>10}')
            for i in range(len(qoi_idx)):
                for j in range(xdim):
                    for k in range(j+1, xdim):
                        q = qoi_idx[i]
                        print(f'{qoi_labels[i]:>10} {"("+param_labels[j]+", "+param_labels[k]+")":>20} '
                              f'{S2_est[j, k, q]: 10.3f} {S2_se[j, k, q]: 10.3f}')

        S1_total = np.sum(S1_est, axis=0)       # (ydim,)
        S2_total = np.zeros((ydim,))            # (ydim,)
        if compute_s2:
            for i in range(xdim):
                for j in range(i+1, xdim):
                    S2_total += S2_est[i, j, :]     # sum the upper diagonal
        print(f'\n{"QoI":>10} {"S1 total":>10} {"S2 total":>10} {"Higher order":>15}')
        for i in range(len(qoi_idx)):
            q = qoi_idx[i]
            print(f'{qoi_labels[i]:>10} {S1_total[q]: 10.3f} {S2_total[q]: 10.3f} '
                  f'{1 - S1_total[q] - S2_total[q]: 15.3f}')

    if plot:
        # Plot bar chart of S1, ST
        c = plt.get_cmap(cmap)
        fig, axs = plt.subplots(1, len(qoi_idx))
        for i in range(len(qoi_idx)):
            ax = axs[i] if len(qoi_idx) > 1 else axs
            q = qoi_idx[i]
            z = st.norm.ppf(1 - (1-0.95)/2)  # get z-score from N(0,1), assuming CLT at n>30
            x = np.arange(xdim)
            width = 0.2
            ax.bar(x - width / 2, S1_est[:, q], width, color=c(0.1), yerr=S1_se[:, q] * z,
                   label=r'$S_1$', capsize=3, linewidth=1, edgecolor=[0, 0, 0])
            ax.bar(x + width / 2, ST_est[:, q], width, color=c(0.9), yerr=ST_se[:, q] * z,
                   label=r'$S_{T}$', capsize=3, linewidth=1, edgecolor=[0, 0, 0])
            ax_default(ax, "Model parameters", "Sobol' index", legend=True)
            ax.set_xticks(x, param_labels)
            ax.set_ylim(bottom=0)
            ax.set_title(qoi_labels[i])
        fig.set_size_inches(4*len(qoi_idx), 4)
        fig.tight_layout()
        bar_chart = (fig, axs)

        # Plot pie chart of S1, S2, higher-order
        fig, axs = plt.subplots(1, len(qoi_idx))
        for i in range(len(qoi_idx)):
            ax = axs[i] if len(qoi_idx) > 1 else axs
            q = qoi_idx[i]
            values = []
            labels = []
            s12_other = 0
            thresh = 0.05    # Only show indices with > 5% effect
            for j in range(xdim):
                if S1_est[j, q] > thresh:
                    values.append(S1_est[j, q])
                    labels.append(param_labels[j])
                else:
                    s12_other += max(S1_est[j, q], 0)

            if compute_s2:
                for j in range(xdim):
                    for k in range(j+1, xdim):
                        if S2_est[j, k, q] > thresh:
                            values.append(S2_est[j, k, q])
                            labels.append("("+param_labels[j]+", "+param_labels[k]+")")
                        else:
                            s12_other += max(S2_est[j, k, q], 0)

            values.append(max(s12_other, 0))
            labels.append(r'Other $S_1$, $S_2$')
            s_higher = max(1 - np.sum(values), 0)
            values.append(s_higher)
            labels.append(r'Higher order')

            # Adjust labels to show percents, sort by value, and threshold small values for plotting
            labels = [f"{label}, {100*values[i]:.1f}%" if values[i] > thresh else
                      f"{label}, <{max(0.5, round(100*values[i]))}%" for i, label in enumerate(labels)]
            values = [val if val > thresh else max(0.02, val) for val in values]
            labels, values = list(zip(*sorted(zip(labels, values), reverse=True, key=lambda ele: ele[1])))

            # Generate pie chart
            colors = c(np.linspace(0, 1, len(values)-2))
            gray_idx = [idx for idx, label in enumerate(labels) if label.startswith('Higher') or
                        label.startswith('Other')]
            pie_colors = np.empty((len(values), 4))
            c_idx = 0
            for idx in range(len(values)):
                if idx in gray_idx:
                    pie_colors[idx, :] = [0.7, 0.7, 0.7, 1]
                else:
                    pie_colors[idx, :] = colors[c_idx, :]
                    c_idx += 1
            radius = 2
            wedges, label_boxes = ax.pie(values, colors=pie_colors, radius=radius, startangle=270,
                                         shadow=True, counterclock=False, frame=True,
                                         wedgeprops=dict(linewidth=1.5, width=0.6*radius, edgecolor='w'),
                                         textprops={'color': [0, 0, 0, 1], 'fontsize': 10, 'family': 'serif'})
            kw = dict(arrowprops=dict(arrowstyle="-"), zorder=0, va="center", fontsize=9, family='serif',
                      bbox=dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0))

            # Put annotations with arrows to each wedge (coordinate system is relative to center of pie)
            for j, wed in enumerate(wedges):
                ang = (wed.theta2 - wed.theta1) / 2. + wed.theta1
                x = radius * np.cos(np.deg2rad(ang))
                y = radius * np.sin(np.deg2rad(ang))
                ax.scatter(x, y, s=10, c='k')
                kw["horizontalalignment"] = "right" if int(np.sign(x)) == -1 else "left"
                kw["arrowprops"].update({"connectionstyle": f"angle,angleA=0,angleB={ang}"})
                y_offset = 0.2 if j == len(labels) - 1 else 0
                ax.annotate(labels[j], xy=(x, y), xytext=((radius+0.2)*np.sign(x), 1.3*y - y_offset), **kw)
            ax.set(aspect="equal")
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['left'].set_visible(False)
            ax.spines['bottom'].set_visible(False)
            ax.get_yaxis().set_ticks([])
            ax.get_xaxis().set_ticks([])
            ax.set_title(qoi_labels[i])
        fig.set_size_inches(3*radius*len(qoi_idx), 2.5*radius)
        fig.tight_layout()
        fig.subplots_adjust(left=0.15, right=0.75)
        pie_chart = (fig, axs)

    if compute_s2:
        ret = (S1, S2, ST)
    else:
        ret = (S1, ST)
    if plot:
        ret = ret + (bar_chart, pie_chart)
    return ret