Tutorial 03 - Computation of Sobol Indices

This tutorial covers the approximation of the Sobol indices of a target function, which are used to infer information about the global parameter sensitivity of the model. To verify the results we compute with PyThia, we use the Sobol function as the object of interest, as the Sobol indices are explicitly known for this function. The Sobol function is given by

\[f(x) = \prod_{j=1}^{M} \frac{|4 x_j-2| + a_j}{1 + a_j} \qquad\mbox{for}\quad a_1,\dots,a_M \geq 0.\]

Note

The larger \(a_j\), the less is the influence, i.e. the sensitivity, of parameter \(x_j\).

The mean of the Sobol function is \(\mathbb{E}[f]=1\) and the variance reads

\[\operatorname{Var}[f] = \prod_{j=1}^{M} (1+c_j) - 1 \qquad\mbox{for}\qquad c_j = \frac{1}{3(1+a_j)^{2}}.\]

With this, we can easily compute the Sobol indices by

\[\operatorname{Sob}_{i_1, ..., i_s} = \frac{1}{\operatorname{Var}[f]}\prod_{k=1}^{s} c_{i_k}.\]

An implementation of the Sobol function method sobol_function() and the analytical Sobol indices method sobol_sc() is included in the complete script at the end of this tutorial.

First, we choose some values for the coefficients \(a_1,\dots,a_M\) and with this the target function.

import numpy as np
a = np.array([1, 2, 3])
def target_function(x):
    return sobol_function(x, a=a)

Additionally, we compute the analytical Sobol indices.

sobol_dict = sobol_sc(a=a, dim=len(a))[0]
sobol_coefficients = np.array(list(sobol_dict.values())).reshape(-1, 1)

Next we define the necessary quantities for the PC expansion. For more information see Tutorial 01 - Approximation of Functions with Polynomial Chaos. As stochastic parameters we need to choose uniformly distributed variables \(x_j\) on \([0,1]\) according to the number of coefficients \(a_1,\dots,a_M\), i.e.,

import pythia as pt
params = [pt.parameter.Parameter(
    index=j, name=f"x_{j+1}", domain=[0, 1], distribution='uniform')
    for j in range(a.size)]

As expansion terms we choose multivariate Legendre polynomials of total polynomial degree less then \(11\)

max_dim = 11
# limit total polynomial degree of expansion terms to 10
indices = pt.index.simplex(len(params), max_dim-1)
index_set = pt.index.IndexSet(indices)

The last thing we need are training data. We generate the training pairs again by an optimal weighted distribution, see Tutorial 02 - Approximation of n-D functions with Polynomial Chaos for more detail.

s = pt.sampler.WLSTensorSampler(params, [max_dim-1]*len(params))
x_train = s.sample(10_000)
w_train = s.weight(x_train)
y_train = target_function(x_train)

With this, we compute the PC expansion of the target_function with PyThia

surrogate = pt.chaos.PolynomialChaos(params, index_set, x_train, w_train, y_train)

As the approximative Sobol indices can be easily derived from the PC expansion coefficients, the PolynomialChaos class computes them automatically upon initialization. They can be accessed via surrogate.sobol, which is an array ordered according to index_set.sobol_tuples. Hence it is easy to verify, if the approximation of the Sobol indices was done correctly.

print(f" {'sobol_tuple':<12} {'exact':<8}  {'approx':<8}  {'abs error':<9}")
print("-"*44)
for j, sdx in enumerate(sobol_dict.keys()):
    print(f" {str(sdx):<11} ",  # Sobol index subscripts
          f"{sobol_coefficients[j, 0]:<4.2e} ",
          f"{surrogate.sobol[j, 0]:<4.2e} ",
          f"{np.abs(sobol_coefficients[j, 0] - surrogate.sobol[j, 0]):<4.2e}"
          )

This concludes this tutorial. Below you find the complete script you can use to run on your own system. This script also computes the approximation error of the PC surrogate.

Complete Script

import numpy as np
import pythia as pt


def sobol_function(x, a=None, **kwargs):
    if not 0 < x.ndim < 3:
        raise ValueError('Wrong ndim: {}'.format(x.ndim))
    if x.ndim == 1:
        x.shape = 1, -1
    if a is None:
        a = np.zeros(x.shape[1])
    elif not a.shape == (x.shape[1],):
        raise ValueError('Wrong shape: {}'.format(a.shape))
    return np.prod((abs(4.0*x - 2.0) + a) / (1.0 + a), axis=1).reshape(-1, 1)


def sobol_sc(a, dim=1, **kwargs):
    sobol = {}
    beta = (1.0+a)**(-2) / 3
    var = np.prod(1.0 + beta) - 1.0
    sobol_tuples = pt.index.IndexSet(pt.index.tensor([1, 1, 1])).sobol_tuples
    for sdx in sobol_tuples:
        sobol[sdx] = 1.0 / var
        for k in sdx:
            sobol[sdx] *= beta[k-1]
    if dim > 1:
        return np.array([sobol for _ in range(dim)])
    else:
        return sobol


# target function definition
a = np.array([1, 2, 3])


def target_function(x):
    return sobol_function(x, a=a)


# analytical sobol coefficients
sobol_dict = sobol_sc(a=a, dim=len(a))[0]
sobol_coefficients = np.array(list(sobol_dict.values())).reshape(-1, 1)

# setup pc surrogate
params = [pt.parameter.Parameter(
    index=j, name=f"x_{j+1}", domain=[0, 1], distribution='uniform')
    for j in range(a.size)]

max_dim = 11
# limit total polynomial degree of expansion terms to 10
indices = pt.index.simplex(len(params), max_dim-1)
index_set = pt.index.IndexSet(indices)
print(f"[{pt.misc.now()}] multiindex information:")
print(f"    number of indices: {index_set.shape[0]}")
print(f"    dimension: {index_set.shape[1]}")
print(f"    number of sobol indices: {len(index_set.sobol_tuples)}")

N = 10_000
print(f"[{pt.misc.now()}] generate training data ({N})")
s = pt.sampler.WLSTensorSampler(params, [max_dim-1]*len(params))
x_train = s.sample(N)
w_train = s.weight(x_train)
y_train = target_function(x_train)

print(f"[{pt.misc.now()}] compute pc expansion")
surrogate = pt.chaos.PolynomialChaos(
    params, index_set, x_train, w_train, y_train)

# test PC approximation
N = 1000
print(f"[{pt.misc.now()}] generate test data ({N})")
s_test = pt.sampler.ParameterSampler(params)
x_test = s_test.sample(N)
y_test = target_function(x_test)
y_approx = surrogate.eval(x_test)

error_L2 = np.sqrt(np.sum((y_test-y_approx)**2)/N)
error_L2_rel = error_L2 / np.sqrt(np.sum((y_test)**2)/N)
error_max = np.max(np.abs(y_test-y_approx))
error_max_rel = np.max(np.abs(y_test-y_approx)/np.abs(y_test))

print(f"[{pt.misc.now()}]     test error L2 (abs/rel):",
      f" {error_L2:4.2e} / {error_L2_rel:4.2e}")
print(f"[{pt.misc.now()}]     test error max (abs/rel):",
      f" {error_max:4.2e} / {error_max_rel:4.2e}")

# compare Sobol indices
print(f"[{pt.misc.now()}] Comparison of Sobol indices")
print(f" {'sobol_tuple':<12} {'exact':<8}  {'approx':<8}  {'abs error':<9}")
print("-"*44)
for j, sdx in enumerate(sobol_dict.keys()):
    print(f" {str(sdx):<11} ",  # Sobol index subscripts
          f"{sobol_coefficients[j, 0]:<4.2e} ",
          f"{surrogate.sobol[j, 0]:<4.2e} ",
          f"{np.abs(sobol_coefficients[j, 0] - surrogate.sobol[j, 0]):<4.2e}"
          )