Example – Surrogate for stochastic heat equation with principal component analysis

Example – Surrogate for stochastic heat equation with principal component analysis#

In the previous example, we built a surrogate for the solution to the stochastic (steady-state) heat equation at a single point in space. In this notebook, we will build a surrogate for the entire temperature field.

We will reuse the same probability measure on conductivity \(c\), i.e.,

\[ c(x; \xi) = c_0(x) \exp\Big\{ g(x; \xi) \Big\} \]

where \(g\) is a Gaussian process. As before, we implement it with the Karhunen-Loève expansion:

Show code cell source Hide code cell source

class KarhunenLoeveExpansion(object):
    
    """
    A class representing the Karhunen Loeve Expansion of a Gaussian random field.
    It uses the Nystrom approximation to do it.
    
    Arguments:
        k      -     The covariance function.
        Xq     -     Quadrature points for the Nystrom approximation.
        wq     -     Quadrature weights for the Nystrom approximation.
        alpha  -     The percentage of the energy of the field that you want to keep.
        X      -     Observed inputs (optional).
        Y      -     Observed field values (optional).
    """
    
    def __init__(self, k, Xq=None, wq=None, nq=100, alpha=0.9, X=None, y=None, *, input_dim):

        self.k = k
        if input_dim is None:
            self.input_dim = 1
        else:
            self.input_dim = input_dim

        # Generate quadrature points
        if Xq is None:
            if input_dim is None:
                self.input_dim = 1
            if self.input_dim == 1:
                Xq = jnp.linspace(0, 1, nq)[:, None]
                wq = jnp.ones((nq, )) / nq
            elif self.input_dim == 2:
                nq = int(jnp.sqrt(nq))
                x = jnp.linspace(0, 1, nq)
                X1, X2 = jnp.meshgrid(x, x)
                Xq = jnp.hstack([X1.flatten()[:, None], X2.flatten()[:, None]])
                wq = jnp.ones((nq ** 2, )) / nq ** 2
            else:
                raise NotImplementedError('For more than 2D, please supply quadrature points and weights.')
        else:
            self.input_dim = Xq.shape[1]
        self.Xq = Xq
        self.wq = wq
        self.k = k
        self.alpha = alpha

        # Evaluate the covariance function at the quadrature points
        # If we have some observed data, we need to use the posterior covariance
        if X is not None:
            self.D = gpx.Dataset(X, y[:, None])
            prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=k)
            likelihood = gpx.likelihoods.Gaussian(num_datapoints=X.shape[0])
            posterior = prior * likelihood
            posterior.likelihood.obs_stddev = gpx.parameters.Static(1e-6)
            self.posterior = posterior
            Kq = posterior.predict(Xq, train_data=self.D).covariance()
        else:
            self.D = None
            self.prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=k)
            Kq = self.prior.predict(Xq).covariance()
        
        # Get the eigenvalues/eigenvectors of the discretized covariance function
        B = jnp.einsum('ij,j->ij', Kq, wq)
        lam, v = jax.scipy.linalg.eigh(B, overwrite_a=True)
        lam = lam[::-1]
        lam = lam.at[lam <= 0.].set(0.)

        # Keep only the eigenvalues that explain alpha% of the energy
        energy = jnp.cumsum(lam) / jnp.sum(lam)
        i_end = jnp.arange(energy.shape[0])[energy > alpha][0] + 1
        lam = lam[:i_end]
        v = v[:, ::-1]
        v = v[:, :i_end]
        
        self.lam = lam
        self.sqrt_lam = jnp.sqrt(lam)
        self.v = v
        self.energy = energy
        self.num_xi = i_end
        
    def eval_phi(self, x):
        """
        Evaluate the eigenfunctions at x.
        """
        if self.D is not None:
            nq = self.Xq.shape[0]
            Xf = jnp.vstack([self.Xq, x])
            latent_dist = self.posterior.predict(Xf, train_data=self.D)
            m = latent_dist.mean()
            C = latent_dist.covariance()
            Kc = C[:nq, nq:].T
            self.tmp_mu = m[nq:]
        else:
            Kc = self.prior.kernel.cross_covariance(x, self.Xq)
            self.tmp_mu = 0.
        phi = jnp.einsum("i,ji,j,rj->ri", 1. / self.lam, self.v, self.wq**0.5, Kc)
        return phi
    
    def __call__(self, x, xi):
        """
        Evaluate the expansion at x and xi.
        """
        phi = self.eval_phi(x)
        return self.tmp_mu + jnp.dot(phi, xi * self.sqrt_lam)

k = gpx.kernels.RBF(lengthscale=0.1, variance=0.5)
kle = KarhunenLoeveExpansion(k, nq=1000, alpha=0.95, input_dim=1)

def c(x, xi):
    """Compute the random thermal conductivity field for a given xi."""
    return jnp.exp(kle(x, xi))

../../_images/d03d371cb38bd704680cde3c26afbe13aa5140caa8ea47271df05f5eb39ee2b5.svg

Reduce dimensionality of the output with PCA#

The goal is to propagate uncertainty from one stochastic function (the conductivity field) to another stochastic function (the temperature field). This requires approximating the infinite-dimensional fields with something finite-dimensional.

We have already represented the input field with a finite number of Karhunen-Loève expansion coefficients (as done in the previous example). In a similar vein, we will represent the output temperature field with a finite number of principal component analysis (PCA) coefficients.

Let’s do it. First we create a solver for the heat equation:

solver = SteadyStateHeat1DSolver(nx=500)

Next, create some training (and testing) data:

xis = []
us = []
for i in range(100):
    key, key_xi = jrandom.split(key)
    xi = jrandom.normal(key_xi, shape=(kle.num_xi,))
    x, y = solver(xi)
    xis.append(xi)
    us.append(y)
xis = jnp.stack(xis, axis=0)
us = jnp.stack(us, axis=0)

from sklearn.model_selection import train_test_split
key, key_split = jrandom.split(key)
xi_train, xi_test, u_train, u_test = train_test_split(xis, us, test_size=0.1, random_state=int(jrandom.randint(key_split, shape=(), minval=0, maxval=1e6)))

And finally, do PCA on the training data:

class PrincipalComponentAnalysis(object):
    def __init__(self, X, alpha=0.9):
        self.x_bar = jnp.mean(X, axis=0)
        B = X - self.x_bar
        U, s, Vt = jnp.linalg.svd(B, full_matrices=False)
        self.U = U
        self.S = s
        self.Vt = Vt
        energy = jnp.cumsum(s)/jnp.sum(s)
        self.n_components = jnp.arange(energy.shape[0])[energy > alpha][0] + 1
        self._coefficients = jnp.einsum('ij,j->ij', U, s)
    
    @property
    def coefficients(self, n_components=None):
        if n_components is None:
            n_components = self.n_components
        return self._coefficients[:, :n_components]
    
    @property
    def principal_components(self):
        return self.Vt

    def project(self, X, n_components=None):
        if n_components is None:
            n_components = self.n_components
        return jnp.einsum('ij,jk->ik', X, self.Vt[:n_components]) + self.x_bar

pca = PrincipalComponentAnalysis(u_train, alpha=0.95)

Both the input and output of our stochastic model can now be expressed as a weighted sum of fixed functions. Let’s visualize what these fixed functions are:

../../_images/08413835ff4d432de91e00e41493bb7c3ed98185545a1a05112581d0cf9065c1.svg

To reiterate, every conductivity field can be approximately represented as the weighted sum

\[ c(x) \approx \sum_{i=1}^{N_\text{KLE}} \xi_i \sqrt{\lambda_i} \phi^\text{KLE}_i(x) \]

where the fixed functions \(\{\phi^\text{KLE}_i | i=1, \dots, N_\text{KLE}\}\) are shown on the left in the figure above.

Likewise, every temperature field can be approximately represented as the weighted sum

\[ u(x) \approx \sum_{i=1}^{N_\text{PCA}} \eta_i \phi^\text{PCA}_i(x) \]

where the fixed functions \(\{\phi^\text{PCA}_i | i=1, \dots, N_\text{PCA}\}\) are shown on the right in the figure above.

Surrogate for the stochastic model#

We will now train a surrogate to approximate the mapping \(\xi \mapsto \eta\) (i.e., the mapping from the input’s KLE coefficients to the output’s PCA coefficients). This will enable fast uncertainty propagation over fields!

However, instead of directly creating one big multi-input-multi-output surrogate, we will train a separate Gaussian process surrogate for each PCA coefficient. So, there will be a GP for the mapping \(\xi \mapsto \eta_1\), another GP for \(\xi \mapsto \eta_2\), and another for \(\xi \mapsto \eta_3\), and so on.

Here we go:

pca_datasets = []
posteriors = []
for i in range(pca.n_components):
    key, key_train = jrandom.split(key)
    D = gpx.Dataset(xi_train, pca.coefficients[:, i][:, None])
    posterior = train_gp(D, measurement_noise=1e-6, key=key_train, verbose=True)
    pca_datasets.append(D)
    posteriors.append(posterior)

Great! We now have a trained GP for each PCA coefficient of the temperature field. Let’s group it all together in a single class for convenience:

surrogate = Surrogate(posteriors, pca_datasets, x, pca)

And let’s test our surrogate’s accuracy. First, let’s predict the full temperature field for each test case:

u_test_pred = surrogate(xi_test)

And let’s compare them to the true fields with a parity plot:

../../_images/c835f57a8ebbdd4d3cff53786b07738d2f34bfa4128d70f89b23419f61487664.svg

The fit looks good.

Uncertainty propagation#

We can now use the surrogate to do uncertainty quantification tasks for very cheap!

For example, let’s compute some quantiles for the distribution of \(u(x)\):

n_samples = 1000
key, key_xi = jrandom.split(key)
xi = jrandom.normal(key_xi, shape=(n_samples, kle.num_xi))
u = surrogate(xi)

u_mean = jnp.mean(u, axis=0)
u_025, u_975 = jnp.quantile(u, jnp.array([0.025, 0.975]), axis=0)

../../_images/46d93848b06db47c1ea4bf08aa9413804278c6b021743c3fa5df412ce45759f3.svg

Take a moment to appreciate this. With our newly built surrogate, we have propagated uncertainty through a partial differential equation from an infinite dimensional random field to another infinite dimensional random field. And evaluating the distribution of the entire output field only takes a fraction of a second!

Surrogates for stochastic models (like this one) can enable downstream tasks that were prohibitively expensive before. For example, if we include some design variables in the input, we could perform design optimization under uncertainty.

Questions#

Play with the amount of training data, as well as number of KLE and PCA coefficients (controlled by the alpha parameters).