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

a = np.array([1, 2, 3])
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)
]

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_set(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.

N = 10_000
s = pt.least_squares_sampler.OWLSTensorSampler(params, [max_dim - 1] * len(params))
x_train = s.sample(N)
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("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}",
    )

This concludes this tutorial.

Complete Script

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

  1"""Tutorial 03 - Computation of Sobol Indices."""
  2
  3from typing import Any
  4
  5import numpy as np
  6import pythia as pt
  7
  8
  9def sobol_function(
 10    x: np.ndarray, a: np.ndarray | None = None, **kwargs: Any
 11) -> np.ndarray:
 12    """Sobol function.
 13
 14    Parameters
 15    ----------
 16    x : np.ndarray
 17        Input values.
 18    a : np.ndarray | None, optional
 19        Coefficients, by default None
 20
 21    Returns
 22    -------
 23    :
 24        Sobol function values.
 25
 26    Raises
 27    ------
 28    ValueError
 29        Wrong dimension for input `x`.
 30    ValueError
 31        Wrong shape of Coefficients `a`.
 32    """
 33    if not 0 < x.ndim < 3:
 34        raise ValueError("Wrong ndim: {}".format(x.ndim))
 35    if x.ndim == 1:
 36        x.shape = 1, -1
 37    if a is None:
 38        a = np.zeros(x.shape[1])
 39    elif not a.shape == (x.shape[1],):
 40        raise ValueError("Wrong shape: {}".format(a.shape))
 41    return np.prod((abs(4.0 * x - 2.0) + a) / (1.0 + a), axis=1).reshape(-1, 1)
 42
 43
 44def sobol_sc(a: np.ndarray, dim: int = 1, **kwargs: Any) -> dict | np.ndarray:
 45    """Sobol function Sobol indices.
 46
 47    Parameters
 48    ----------
 49    a : np.ndarray
 50        Coefficients.
 51    dim : int, optional
 52        Parameter dimension, by default 1.
 53
 54    Returns
 55    -------
 56    :
 57        Sobol indices of Sobol function.
 58    """
 59    sobol = {}
 60    beta = (1.0 + a) ** (-2) / 3
 61    var = np.prod(1.0 + beta) - 1.0
 62    sobol_tuples = pt.index.IndexSet(pt.index.tensor_set([1, 1, 1])).sobol_tuples
 63    for sdx in sobol_tuples:
 64        sobol[sdx] = 1.0 / var
 65        for k in sdx:
 66            sobol[sdx] *= beta[k - 1]
 67    if dim > 1:
 68        return np.array([sobol for _ in range(dim)])
 69    else:
 70        return sobol
 71
 72
 73print("Tutorial 03 - 2D approximation with PC")
 74
 75# target function definition
 76a = np.array([1, 2, 3])
 77
 78
 79def target_function(x: np.ndarray) -> np.ndarray:
 80    """Target function.
 81
 82    Parameters
 83    ----------
 84    x : np.ndarray
 85    """
 86    return sobol_function(x, a=a)
 87
 88
 89# analytical sobol coefficients
 90sobol_dict = sobol_sc(a=a, dim=len(a))[0]
 91sobol_coefficients = np.array(list(sobol_dict.values())).reshape(-1, 1)
 92
 93# setup pc surrogate
 94params = [
 95    pt.parameter.Parameter(name=f"x_{j+1}", domain=(0, 1), distribution="uniform")
 96    for j in range(a.size)
 97]
 98
 99max_dim = 11
100# limit total polynomial degree of expansion terms to 10
101indices = pt.index.simplex_set(len(params), max_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.least_squares_sampler.OWLSTensorSampler(params, [max_dim - 1] * 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 = pt.chaos.PolynomialChaos(params, index_set, x_train, w_train, y_train)
117
118# test PC approximation
119N = 1000
120print(f"generate test data ({N})")
121s_test = pt.sampler.ParameterSampler(params)
122x_test = s_test.sample(N)
123y_test = target_function(x_test)
124y_approx = surrogate.eval(x_test)
125
126error_L2 = np.sqrt(np.sum((y_test - y_approx) ** 2) / N)
127error_L2_rel = error_L2 / np.sqrt(np.sum((y_test) ** 2) / N)
128error_max = np.max(np.abs(y_test - y_approx))
129error_max_rel = np.max(np.abs(y_test - y_approx) / np.abs(y_test))
130
131print(f"test error L2 (abs/rel): {error_L2:4.2e} / {error_L2_rel:4.2e}")
132print(f"test error max (abs/rel): {error_max:4.2e} / {error_max_rel:4.2e}")
133
134# compare Sobol indices
135print("Comparison of Sobol indices")
136print(f" {'sobol_tuple':<12} {'exact':<8}  {'approx':<8}  {'abs error':<9}")
137print("-" * 44)
138for j, sdx in enumerate(sobol_dict.keys()):
139    print(
140        f" {str(sdx):<11} ",  # Sobol index subscripts
141        f"{sobol_coefficients[j, 0]:<4.2e} ",
142        f"{surrogate.sobol[j, 0]:<4.2e} ",
143        f"{np.abs(sobol_coefficients[j, 0] - surrogate.sobol[j, 0]):<4.2e}",
144    )