Tutorial 04 - Automatic choice of expansion terms

This tutorial covers the automatic generation of sparse PC expansion indices based on a crude approximation of the Sobol indices. 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, 4, 8])
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)]

We want to compare the approximation results of the PC expansion using automatically generated expansion indices with a choice done by hand. For this we compute a reference PC expansion using multivariate Legendre polynomials of total degree less then \(7\).

dim = 7
indices = pt.index.simplex(len(params), 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, [dim]*len(params))
x_train = s.sample(10_000)
w_train = s.weight(x_train)
y_train = target_function(x_train)

With this, we can compute a reference PC expansion for our forward model with

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

With this out of the way we now focus on choosing a more sparse set of PC expansion indices automatically. To fit our needs, we specify the maximal number of expansion terms we want the PC expansion to have as well as the truncation threshold. Computing the “optimal” multiindices can then be done using pt.chaos.find_optimal_indices. To show the efficiency of the automatic multiindex computation, we choose an expansion with half the number of terms (\(\sim 100\)) as the reference PC (\(\sim 200\)) and set the truncation threshold to \(10^{-3}\).

max_terms = index_set.shape[0]//2
threshold = 1.0e-03
indices, sC = pt.chaos.find_optimal_indices(
    params, x_train, w_train, y_train, max_terms=max_terms, threshold=threshold)
index_set = pt.index.IndexSet(indices)

For a detailed explaination on the workings of pt.chaos.find_optimal_indices we refer to the module documentation. However, we want to provide a small explaination of the threshold input. In principle find_optimal_indices computes a very inaccurate PC expansion with far too many terms for the given amount of training data to obtain a very crude approximation of as many Sobol indices as possible. The expansion is chosen so that at least one expansion term is computed for each Sobol index. Depending on the threshold the function decides which Sobol indices, i.e., parameter combinations are most relevant. This means, that the function will exclude all multiindices associated to Sobol indices that have a (combined) contribution of less then threshold. The multiindices are then distributed according to the magnitude of the crude Sobol index computation. We use this option here only for showcasing the results later on.

Computing the PC expansion with the automatically chosen multiindices can now be done as before.

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

What remains is to check if the approximation of the Sobol function using only half the expansion terms is comparable in both the approximation error as well as the approximate Sobol indices. If everything went right you should see that the approximation error of both PC expansions is of the same order of magnitude and that the Sobol indices coincide where they are greater then \(10^{-3}\).

error_L2_ref = np.sqrt(np.sum((f_test-f_approx_ref)**2)/x_test.shape[0])
error_L2 = np.sqrt(np.sum((f_test-f_approx)**2)/x_test.shape[0])
print(f"[{pt.misc.now()}]     test error L2 auto: {error_L2:4.2e}")
print(f"[{pt.misc.now()}]     test error L2 ref:  {error_L2_ref:4.2e}")

print(f"[{pt.misc.now()}] Comparison of Sobol indices")
print(f" {'sobol_tuples':<13} {'exact':<8}  {'pc-auto':<8}  {'(crude)':<8}  {'pc-ref':<8}")
print("-"*54)
for j, sdx in enumerate(index_set.sobol_tuples):
    print(f" {str(sdx):<12} ",  # Sobol index subscripts
          f"{sobol_coefficients[j, 0]:<4.2e} ",
          f"{surrogate.sobol[j, 0]:<4.2e} ",
          f"{sC[j][0]:<4.2e} ",
          f"{surrogate_ref.sobol[j, 0]:<4.2e} ")

This concludes this tutorial. Below you find the complete script you can use to run on your own system.

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]*len(a))).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

# function definitions
a = np.array([1, 2, 4, 8])
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 reference PC surrogate
params = [pt.parameter.Parameter(
    index=j, name=f"x_{j+1}", domain=[0, 1], distribution='uniform')
    for j in range(a.size)]

dim = 7
indices = pt.index.simplex(len(params), 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, [dim]*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_ref = pt.chaos.PolynomialChaos(params, index_set, x_train, w_train, y_train)

# auto-generate PC multiindices
max_terms = index_set.shape[0]//2
threshold = 1.0e-03
print(f"[{pt.misc.now()}] compute optimal mdx")
indices, sC = pt.chaos.find_optimal_indices(
    params, x_train, w_train, y_train, max_terms=max_terms, threshold=threshold)
index_set = pt.index.IndexSet(indices)
print(f"[{pt.misc.now()}] automatic 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)}")

print(f"[{pt.misc.now()}] compute optimal 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)
f_test = target_function(x_test)
f_approx = surrogate.eval(x_test)
f_approx_ref = surrogate_ref.eval(x_test)

error_L2_ref = np.sqrt(np.sum((f_test-f_approx_ref)**2)/x_test.shape[0])
error_L2 = np.sqrt(np.sum((f_test-f_approx)**2)/x_test.shape[0])
print(f"[{pt.misc.now()}]     test error L2 auto: {error_L2:4.2e}")
print(f"[{pt.misc.now()}]     test error L2 ref:  {error_L2_ref:4.2e}")

# print Sobol coefficients
print(f"[{pt.misc.now()}] Comparison of Sobol indices")
print(f" {'sobol_tuples':<13} {'exact':<8}  {'pc-auto':<8}  {'(crude)':<8}  {'pc-ref':<8}")
print("-"*54)
for j, sdx in enumerate(index_set.sobol_tuples):
    print(f" {str(sdx):<12} ",  # Sobol index subscripts
          f"{sobol_coefficients[j, 0]:<4.2e} ",
          f"{surrogate.sobol[j, 0]:<4.2e} ",
          f"{sC[j][0]:<4.2e} ",
          f"{surrogate_ref.sobol[j, 0]:<4.2e} ",
          )