The OptionsManager and the AppContext

The PETSc options provide a simple but powerful DSL for configuring composable solvers. However, their main limitation is that the values of each option is limited to primitive C types, e.g. str, float, int, or complex. Sometimes more advanced data is useful or essential for building a particular solver.

petsctools.AppContext fulfils this need by providing a means of passing arbitrary Python types through to Python PETSc types (e.g. Python type PCs). In this demo we show how to use the AppContext to pass data to a custom Python type PC using the variable coefficient diffusion equation as an example.

Diffusion equation with variable coefficients

The diffusion equation with coefficient \(\sigma(x)\) depending on the spatial coordinate is:

\[u - \nabla\cdot\left(\sigma(x)\nabla u\right) = b\]

We will solve this matrix with finite differences with the standard 3 point central stencil. The particular details of the discretisation are not essential for this demo so we will be brief in the description.

If \(D\) is the assembled matrix for the finite difference gradient stencil, and \(\Sigma\) is a diagonal matrix with the value of the diffusion coefficient at each grid point, then the assembled matrix for the diffusion equation is:

\[Au = \left(I + D^{T}\Sigma D\right)u = b, \quad \Sigma_{ii} = \sigma(x_{i})\]

The following Python function takes a numpy array sigma with the value of \(\sigma\) at each grid point and assembles a sparse (aij) PETSc Mat for the diffusion equation. We will use it later to build the Mat for a KSP to solve the diffusion equation.

import numpy as np
import petsctools


def diffusion_mat(sigma):
    """
    AIJ Mat for the diffusion equation with a variable diffusion coefficient.

    (I + D.T@sigma@D)u = b
    """
    from petsc4py import PETSc
    n = sigma.shape[0]
    dtype = sigma.dtype

    # index lists for CSR format
    row_start = [0]
    col_indices = []

    # top row
    idxs = [0, 1]
    col_indices.extend(idxs)
    row_start.append(row_start[-1] + len(idxs))

    # interior rows
    for j in range(1, n-1):
        idxs = [j-1, j, j+1]
        col_indices.extend(idxs)
        row_start.append(row_start[-1] + len(idxs))

    # bottom row
    idxs = [n-2, n-1]
    col_indices.extend(idxs)
    row_start.append(row_start[-1] + len(idxs))

    # values for leading and upper/lower diagonals
    diagonal = 1 + sigma.copy()
    diagonal[:-1] += sigma[1:]
    offdiags = -sigma[1:]

    # interleave diagonal entries
    Avals = np.zeros(3*n-2, dtype=dtype)
    Avals[::3] = diagonal
    Avals[1::3] = offdiags
    Avals[2::3] = offdiags

    A = PETSc.Mat().createAIJWithArrays(
        size=(n, n),
        csr=(row_start, col_indices, Avals)
    )
    return A

A PC needing Python data

Imagine a scenario where we need to solve \(A\) multiple times and \(\sigma\) changes slightly each time, for example if we are solving the unsteady diffusion equation with time-varying coefficients. Rather than recomputing a preconditioner every time \(A\) changes, we might instead find a representative \(\sigma_{p}\) and use that to compute a preconditioner which can be reused for all solves.

For simplicity, in this demo the preconditioner \(P\) will just be the diagonal of the assembled matrix \(A_{p}\) for the diffusion equation with \(\sigma_{p}\) with a simple scaling factor \(\omega\):

\[A_{p} = I + D^{T}\Sigma_{p} D, \quad P = \omega^{-1}\mathrm{diag}(A_{p}).\]

The diagonal of a matrix is clearly not expensive to compute, but in practice we would use a factorisation of \(A_{p}\) which would be more expensive to compute and so would be more worthwhile reusing.

The preconditioner defined above is implemented with a Python type PC called DiffusionJacobiPC in the code below. Constructing \(P\) requires two values, \(\sigma_{p}\) and \(\omega\), which must be provided by the user.

1. The scaling factor \(\omega\) is just a real number, and can therefore be passed as usual via the PETSc.Options using the "djacobi_scale" option.

2. The diffusion coefficient at each grid point \(\sigma_{p}(x_{i})\) is defined as a numpy array. This is clearly not a primitive type and so cannot be passed via the PETSc.Options directly. Instead, we access it via the AppContext using the "djacobi_sigma" key.

The AppContext mimics the PETSc.Options very closely, but can contain arbitrary Python data. We will see below how to add sigma into the AppContext so that it is available to the DiffusionJacobiPC.

class DiffusionJacobiPC:
    prefix = "djacobi_"

    def setFromOptions(self, pc):
        from petsc4py import PETSc
        prefix = (pc.getOptionsPrefix() or "") + self.prefix

        options = PETSc.Options()
        scale = options.getReal(prefix + "scale", 1.0)

        appctx = petsctools.AppContext()
        sigma = appctx[prefix + "sigma"]

        Ap = diffusion_mat(sigma)
        P = Ap.getDiagonal()
        P.scale(1/scale)
        self.P = P

    def apply(self, pc, x, y):
        y.pointwiseDivide(x, self.P)

Assembling the system

We specify the diffusion coefficient as some random variations \(\sigma'\) around a constant value \(\overline{\sigma}\), i.e. \(\sigma(x) = \overline{\sigma} + \sigma'(x)\). Assuming that \(\sigma'\) is the component that may vary from solve to solve, we use \(\sigma_{p}=\overline{\sigma}\).

PETSc = petsctools.init()
np.random.seed(13)
n = 50

sigma_bar = 2*np.ones(n)
sigma_prime = -0.2 + 0.4*np.random.random_sample(n)

sigma = sigma_bar + sigma_prime
sigma_p = sigma_bar

A = diffusion_mat(sigma)

ksp = PETSc.KSP().create()
ksp.setOperators(A)

The Options and the AppContext

Now we configure ksp by passing PETSc options as key-value pairs in the parameters dictionary to petsctools.set_from_options. This function will create a petsctools.OptionsManager and attach it to ksp.

We can see that common options, e.g. "ksp_type" are set as usual in the parameters dictionary. However, when we come to the "djacobi_sigma" value we use the petsctools.AppContextManager class. The AppContextManager.add method returns a unique value which is used to associate a PETSc option to whatever data was passed to add. For example, adding the key-value pair "djacobi_sigma": appmngr.add(sigma_p) to the parameters dictionary means that, during the solve, we will be able to access sigma_p via AppContext()["djacobi_sigma"].

We then pass the AppContextManager to set_from_options() so that it can be attached to ksp and its data can be made available later on during the solve.

appmngr = petsctools.AppContextManager()

petsctools.set_from_options(
    ksp,
    parameters={
        'ksp_converged_reason': None,
        'ksp_type': 'richardson',
        'pc_type': 'python',
        'pc_python_type': f'{__name__}.DiffusionJacobiPC',
        'djacobi_scale': 0.9,
        'djacobi_sigma': appmngr.add(sigma_p),
    },
    appmngr=appmngr,
    options_prefix="",
)

Solving the KSP

Now we come to actually solving the linear equation \(Au=b\). To avoid memory leaks, set_from_options() does not permanently insert the contents of parameters and the appmngr into the global PETSc.Options and AppContext databases respectively. Instead, we use the petsctools.inserted_options context manager. On entry, this context manager inserts the contents of parameters and appmngr into the global databases, and on exit it removes them again. This means that we need to use the inserted_options() context manager whenever these entries will be needed, for example during the solve when the KSP and PC are being set up.

u, b = A.createVecs()
u.zeroEntries()
b.array[:] = np.random.random_sample(n)

with petsctools.inserted_options(ksp):
    ksp.solve(b, u)