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
import pythia as pt

a = np.array([1, 2, 4, 8])
def target_function(x: np.ndarray) -> np.ndarray:
    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.,

params = [
    pt.parameter.Parameter(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_set(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.

N = 10_000
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)

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"test error L2 auto: {error_L2:4.2e}")
print(f"test error L2 ref:  {error_L2_ref:4.2e}")

# print Sobol coefficients
print("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.

Complete Script

The complete source code listed below is located in the tutorials/ subdirectory of the repository root.

  1"""
  2File: tutorials/tutorial_03.py
  3Author: Nando Hegemann
  4Gitlab: https://gitlab.com/Nando-Hegemann
  5Description: Automatic choice of expansion terms.
  6SPDX-License-Identifier: LGPL-3.0-or-later OR Hippocratic-3.0-ECO-MEDIA-MIL
  7"""
  8import numpy as np
  9import pythia as pt
 10
 11
 12def sobol_function(x: np.ndarray, a: np.ndarray | None = None, **kwargs) -> np.ndarray:
 13    """Sobol function.
 14
 15    Parameters
 16    ----------
 17    x : np.ndarray
 18        Input values.
 19    a : np.ndarray | None, optional
 20        Coefficients, by default None
 21
 22    Returns
 23    -------
 24    :
 25        Sobol function values.
 26
 27    Raises
 28    ------
 29    ValueError
 30        Wrong dimension for input `x`.
 31    ValueError
 32        Wrong shape of Coefficients `a`.
 33    """
 34    if not 0 < x.ndim < 3:
 35        raise ValueError("Wrong ndim: {}".format(x.ndim))
 36    if x.ndim == 1:
 37        x.shape = 1, -1
 38    if a is None:
 39        a = np.zeros(x.shape[1])
 40    elif not a.shape == (x.shape[1],):
 41        raise ValueError("Wrong shape: {}".format(a.shape))
 42    return np.prod((abs(4.0 * x - 2.0) + a) / (1.0 + a), axis=1).reshape(-1, 1)
 43
 44
 45def sobol_sc(a: np.ndarray, dim: int = 1, **kwargs) -> tuple | np.ndarray:
 46    """Sobol function Sobol indices.
 47
 48    Parameters
 49    ----------
 50    a : np.ndarray
 51        Coefficients.
 52    dim : int, optional
 53        Parameter dimension, by default 1.
 54
 55    Returns
 56    -------
 57    :
 58        Sobol indices of Sobol function.
 59    """
 60    sobol = {}
 61    beta = (1.0 + a) ** (-2) / 3
 62    var = np.prod(1.0 + beta) - 1.0
 63    sobol_tuples = pt.index.IndexSet(pt.index.tensor_set([1] * len(a))).sobol_tuples
 64    for sdx in sobol_tuples:
 65        sobol[sdx] = 1.0 / var
 66        for k in sdx:
 67            sobol[sdx] *= beta[k - 1]
 68    if dim > 1:
 69        return np.array([sobol for _ in range(dim)])
 70    else:
 71        return sobol
 72
 73
 74print("Tutorial 04 - Automatic choice of expansion terms")
 75
 76# function definitions
 77a = np.array([1, 2, 4, 8])
 78
 79
 80def target_function(x: np.ndarray) -> np.ndarray:
 81    """Target function.
 82
 83    Parameters
 84    ----------
 85    x : np.ndarray
 86    """
 87    return sobol_function(x, a=a)
 88
 89
 90# analytical Sobol coefficients
 91sobol_dict = sobol_sc(a=a, dim=len(a))[0]
 92sobol_coefficients = np.array(list(sobol_dict.values())).reshape(-1, 1)
 93
 94# setup reference PC surrogate
 95params = [
 96    pt.parameter.Parameter(name=f"x_{j+1}", domain=[0, 1], distribution="uniform")
 97    for j in range(a.size)
 98]
 99
100dim = 7
101indices = pt.index.simplex_set(len(params), dim - 1)
102index_set = pt.index.IndexSet(indices)
103print("multiindex information:")
104print(f"    number of indices: {index_set.shape[0]}")
105print(f"    dimension: {index_set.shape[1]}")
106print(f"    number of sobol indices: {len(index_set.sobol_tuples)}")
107
108N = 10_000
109print(f"generate training data ({N})")
110s = pt.sampler.WLSTensorSampler(params, [dim] * len(params))
111x_train = s.sample(N)
112w_train = s.weight(x_train)
113y_train = target_function(x_train)
114
115print("compute pc expansion")
116surrogate_ref = pt.chaos.PolynomialChaos(params, index_set, x_train, w_train, y_train)
117
118# auto-generate PC multiindices
119max_terms = index_set.shape[0] // 2
120threshold = 1.0e-03
121print("compute optimal mdx")
122indices, sC = pt.chaos.find_optimal_indices(
123    params, x_train, w_train, y_train, max_terms=max_terms, threshold=threshold
124)
125index_set = pt.index.IndexSet(indices)
126print("automatic multiindex information:")
127print(f"    number of indices: {index_set.shape[0]}")
128print(f"    dimension: {index_set.shape[1]}")
129print(f"    number of sobol indices: {len(index_set.sobol_tuples)}")
130
131print("compute optimal pc expansion")
132surrogate = pt.chaos.PolynomialChaos(params, index_set, x_train, w_train, y_train)
133
134# test PC approximation
135N = 1000
136print(f"generate test data ({N})")
137s_test = pt.sampler.ParameterSampler(params)
138x_test = s_test.sample(N)
139f_test = target_function(x_test)
140f_approx = surrogate.eval(x_test)
141f_approx_ref = surrogate_ref.eval(x_test)
142
143error_L2_ref = np.sqrt(np.sum((f_test - f_approx_ref) ** 2) / x_test.shape[0])
144error_L2 = np.sqrt(np.sum((f_test - f_approx) ** 2) / x_test.shape[0])
145print(f"test error L2 auto: {error_L2:4.2e}")
146print(f"test error L2 ref:  {error_L2_ref:4.2e}")
147
148# print Sobol coefficients
149print("Comparison of Sobol indices")
150print(
151    f" {'sobol_tuples':<13} {'exact':<8}  {'pc-auto':<8}  {'(crude)':<8}  {'pc-ref':<8}"
152)
153print("-" * 54)
154for j, sdx in enumerate(index_set.sobol_tuples):
155    print(
156        f" {str(sdx):<12} ",  # Sobol index subscripts
157        f"{sobol_coefficients[j, 0]:<4.2e} ",
158        f"{surrogate.sobol[j, 0]:<4.2e} ",
159        f"{sC[j][0]:<4.2e} ",
160        f"{surrogate_ref.sobol[j, 0]:<4.2e} ",
161    )