PETSc Python types#
Here we discuss details about Python-aware PETSc types that can be used within the library.
In particular, we discuss matrices, preconditioners, Krylov solvers, nonlinear solvers and ODE integrators.
The low-level, Cython implementation exposing the Python methods is in src/petsc4py/PETSc/libpetsc4py.pyx.
The scripts used here can be found at demo/python_types.
PETSc Python matrix type#
PETSc provides a convenient way to compute the action of linear operators coded
in Python through the petsc4py.PETSc.Mat.Type.PYTHON type.
In addition to the matrix action, the implementation can expose additional methods for use within the library. A template class for the supported methods is given below.
from petsc4py.typing import Scalar
from petsc4py.PETSc import Mat
from petsc4py.PETSc import Vec
from petsc4py.PETSc import IS
from petsc4py.PETSc import InsertMode
from petsc4py.PETSc import NormType
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by MATPYTHON
class MatPythonProtocol:
def mult(self, A: Mat, x: Vec, y: Vec) -> None:
"""Matrix vector multiplication: y = A @ x."""
...
def multAdd(self, A: Mat, x: Vec, y: Vec, z: Vec) -> None:
"""Matrix vector multiplication: z = A @ x + y."""
...
def multTranspose(self, A: Mat, x: Vec, y: Vec) -> None:
"""Transposed matrix vector multiplication: y = A^T @ x."""
...
def multTransposeAdd(self, A: Mat, x: Vec, y: Vec, z: Vec) -> None:
"""Transposed matrix vector multiplication: z = A^T @ x + y."""
...
def multHermitian(self, A: Mat, x: Vec, y: Vec) -> None:
"""Hermitian matrix vector multiplication: y = A^H @ x."""
...
def multHermitianAdd(self, A: Mat, x: Vec, y: Vec, z: Vec) -> None:
"""Hermitian matrix vector multiplication: z = A^H @ x + y."""
...
def view(self, A: Mat, viewer: Viewer) -> None:
"""View the matrix."""
...
def setFromOptions(self, A: Mat) -> None:
"""Process command line for customization."""
...
def multDiagonalBlock(self, A: Mat, x: Vec, y: Vec) -> None:
"""Perform the on-process matrix vector multiplication."""
...
def createVecs(self, A: Mat) -> tuple[Vec, Vec]:
"""Return tuple of vectors (x,y) suitable for A @ x = y."""
...
def scale(self, A: Mat, s: Scalar) -> None:
"""Scale the matrix by a scalar."""
...
def shift(self, A: Mat, s: Scalar) -> None:
"""Shift the matrix by a scalar."""
...
def createSubMatrix(self, A: Mat, r: IS, c: IS, out: Mat) -> Mat:
"""Return the submatrix corresponding to r rows and c columns.
Matrix out must be reused if not None.
"""
...
def zeroRowsColumns(self, A: Mat, r: IS, diag: Scalar, x: Vec, b: Vec) -> None:
"""Zero rows and columns of the matrix corresponding to the index set r.
Insert diag on the diagonal and modify vectors x and b accordingly if not None.
"""
...
def getDiagonal(self, A: Mat, d: Vec) -> None:
"""Compute the diagonal of the matrix: d = diag(A)."""
...
def setDiagonal(self, A: Mat, d: Vec, im: InsertMode) -> None:
"""Set the diagonal of the matrix."""
...
def missingDiagonal(self, A: Mat, d: Vec, im: InsertMode) -> tuple[bool, int]:
"""Return a flag indicating if the matrix is missing a diagonal entry and the location."""
...
def diagonalScale(self, A: Mat, L: Vec, R: Vec) -> None:
"""Perform left and right diagonal scaling if vectors are not None.
A = diag(L)@A@diag(R).
"""
...
def getDiagonalBlock(self, A: Mat) -> Mat:
"""Return the on-process matrix."""
...
def setUp(self, A: Mat) -> None:
"""Perform the required setup."""
...
def duplicate(self, A: Mat, op: Mat.DuplicateOption) -> Mat:
"""Duplicate the matrix."""
...
def copy(self, A: Mat, B: Mat, op: Mat.Structure) -> None:
"""Copy the matrix: B = A."""
...
def productSetFromOptions(
self, A: Mat, prodtype: str, X: Mat, Y: Mat, Z: Mat
) -> bool:
"""The boolean flag indicating if the matrix supports prodtype."""
...
def productSymbolic(
self, A: Mat, product: Mat, producttype: str, X: Mat, Y: Mat, Z: Mat
) -> None:
"""Perform the symbolic stage of the requested matrix product."""
...
def productNumeric(
self, A: Mat, product: Mat, producttype: str, X: Mat, Y: Mat, Z: Mat
) -> None:
"""Perform the numeric stage of the requested matrix product."""
...
def zeroEntries(self, A: Mat) -> None:
"""Set the matrix to zero."""
...
def norm(self, A: Mat, normtype: NormType) -> float:
"""Compute the norm of the matrix."""
...
def solve(self, A: Mat, y: Vec, x: Vec) -> None:
"""Solve the equation: x = inv(A) y."""
...
def solveAdd(self, A: Mat, y: Vec, z: Vec, x: Vec) -> None:
"""Solve the equation: x = inv(A) y + z."""
...
def solveTranspose(self, A: Mat, y: Vec, x: Vec) -> None:
"""Solve the equation: x = inv(A)^T y."""
...
def solveTransposeAdd(self, A: Mat, y: Vec, z: Vec, x: Vec) -> None:
"""Solve the equation: x = inv(A)^T y + z."""
...
def SOR(
self,
A: Mat,
b: Vec,
omega: float,
sortype: Mat.SORType,
shift: float,
its: int,
lits: int,
x: Vec,
) -> None:
"""Perform SOR iterations."""
...
def conjugate(self, A: Mat) -> None:
"""Perform the conjugation of the matrix: A = conj(A)."""
...
def imagPart(self, A: Mat) -> None:
"""Set real part to zero. A = imag(A)."""
...
def realPart(self, A: Mat) -> None:
"""Set imaginary part to zero. A = real(A)."""
...
In the example below, we create an operator that applies the Laplacian operator
on a two-dimensional grid, and use it to solve the associated linear system.
The default preconditioner in the script is petsc4py.PETSc.PC.Type.JACOBI
which needs to access the diagonal of the matrix.
# ------------------------------------------------------------------------
#
# Poisson problem. This problem is modeled by the partial
# differential equation
#
# -Laplacian(u) = 1, 0 < x,y < 1,
#
# with boundary conditions
#
# u = 0 for x = 0, x = 1, y = 0, y = 1
#
# A finite difference approximation with the usual 5-point stencil
# is used to discretize the boundary value problem to obtain a
# nonlinear system of equations. The problem is solved in a 2D
# rectangular domain, using distributed arrays (DAs) to partition
# the parallel grid.
#
# ------------------------------------------------------------------------
# We first import petsc4py and sys to initialize PETSc
import sys
import petsc4py
petsc4py.init(sys.argv)
# Import the PETSc module
from petsc4py import PETSc
# Here we define a class representing the discretized operator
# This allows us to apply the operator "matrix-free"
class Poisson2D:
def __init__(self, da):
self.da = da
self.localX = da.createLocalVec()
# This is the method that PETSc will look for when applying
# the operator. `X` is the PETSc input vector, `Y` the output vector,
# while `mat` is the PETSc matrix holding the PETSc datastructures.
def mult(self, mat, X, Y):
# Grid sizes
mx, my = self.da.getSizes()
hx, hy = (1.0 / m for m in [mx, my])
# Bounds for the local part of the grid this process owns
(xs, xe), (ys, ye) = self.da.getRanges()
# Map global vector to local vectors
self.da.globalToLocal(X, self.localX)
# We can access the vector data as NumPy arrays
x = self.da.getVecArray(self.localX)
y = self.da.getVecArray(Y)
# Loop on the local grid and compute the local action of the operator
for j in range(ys, ye):
for i in range(xs, xe):
u = x[i, j] # center
u_e = u_w = u_n = u_s = 0
if i > 0:
u_w = x[i - 1, j] # west
if i < mx - 1:
u_e = x[i + 1, j] # east
if j > 0:
u_s = x[i, j - 1] # south
if j < ny - 1:
u_n = x[i, j + 1] # north
u_xx = (-u_e + 2 * u - u_w) * hy / hx
u_yy = (-u_n + 2 * u - u_s) * hx / hy
y[i, j] = u_xx + u_yy
# This is the method that PETSc will look for when the diagonal of the matrix is needed.
def getDiagonal(self, mat, D):
mx, my = self.da.getSizes()
hx, hy = (1.0 / m for m in [mx, my])
(xs, xe), (ys, ye) = self.da.getRanges()
d = self.da.getVecArray(D)
# Loop on the local grid and compute the diagonal
for j in range(ys, ye):
for i in range(xs, xe):
d[i, j] = 2 * hy / hx + 2 * hx / hy
# The class can contain other methods that PETSc won't use
def formRHS(self, B):
b = self.da.getVecArray(B)
mx, my = self.da.getSizes()
hx, hy = (1.0 / m for m in [mx, my])
(xs, xe), (ys, ye) = self.da.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
b[i, j] = 1 * hx * hy
# Access the option database and read options from the command line
OptDB = PETSc.Options()
nx, ny = OptDB.getIntArray(
'grid', (16, 16)
) # Read `-grid <int,int>`, defaults to 16,16
# Create the distributed memory implementation for structured grid
da = PETSc.DMDA().create([nx, ny], stencil_width=1)
# Create vectors to hold the solution and the right-hand side
x = da.createGlobalVec()
b = da.createGlobalVec()
# Instantiate an object of our Poisson2D class
pde = Poisson2D(da)
# Create a PETSc matrix of type Python using `pde` as context
A = PETSc.Mat().create(comm=da.comm)
A.setSizes([x.getSizes(), b.getSizes()])
A.setType(PETSc.Mat.Type.PYTHON)
A.setPythonContext(pde)
A.setUp()
# Create a Conjugate Gradient Krylov solver
ksp = PETSc.KSP().create()
ksp.setType(PETSc.KSP.Type.CG)
# Use diagonal preconditioning
ksp.getPC().setType(PETSc.PC.Type.JACOBI)
# Allow command-line customization
ksp.setFromOptions()
# Assemble right-hand side and solve the linear system
pde.formRHS(b)
ksp.setOperators(A)
ksp.solve(b, x)
# Here we programmatically visualize the solution
if OptDB.getBool('plot', True):
# Modify the option database: keep the X window open for 1 second
OptDB['draw_pause'] = 1
# Obtain a viewer of type DRAW
draw = PETSc.Viewer.DRAW(x.comm)
# View the vector in the X window
draw(x)
# We can also visualize the solution by command line options
# For example, we can dump a VTK file with:
#
# $ python poisson2d.py -plot 0 -view_solution vtk:sol.vts:
#
# or obtain the same visualization as programmatically done above as:
#
# $ python poisson2d.py -plot 0 -view_solution draw -draw_pause 1
#
x.viewFromOptions('-view_solution')
PETSc Python preconditioner type#
The protocol for the petsc4py.PETSc.PC.Type.PYTHON preconditioner is:
from petsc4py.PETSc import KSP
from petsc4py.PETSc import PC
from petsc4py.PETSc import Mat
from petsc4py.PETSc import Vec
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by PCPYTHON
class PCPythonProtocol:
def apply(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the preconditioner on vector b, return in x."""
...
def applySymmetricLeft(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the symmetric left part of the preconditioner on vector b, return in x."""
...
def applySymmetricRight(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the symmetric right part of the preconditioner on vector b, return in x."""
...
def applyTranspose(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the transposed preconditioner on vector b, return in x."""
...
def applyMat(self, pc: PC, B: Mat, X: Mat) -> None:
"""Apply the preconditioner on a block of right-hand sides B, return in X."""
...
def preSolve(self, pc: PC, ksp: KSP, b: Vec, x: Vec) -> None:
"""Callback called at the beginning of a Krylov method.
This method is allowed to modify the right-hand side b and the initial guess x.
"""
...
def postSolve(self, pc: PC, ksp: KSP, b: Vec, x: Vec) -> None:
"""Callback called at the end of a Krylov method.
This method is allowed to modify the right-hand side b and the solution x.
"""
def view(self, pc: PC, viewer: Viewer) -> None:
"""View the preconditioner."""
...
def setFromOptions(self, pc: PC) -> None:
"""Process command line for customization."""
...
def setUp(self, pc: PC) -> None:
"""Perform the required setup."""
...
def reset(self, pc: PC) -> None:
"""Reset the preconditioner."""
...
In the example below, we create a Jacobi preconditioner, which needs to access the diagonal of the matrix. The action of the preconditioner consists of the pointwise multiplication of the inverse diagonal with the input vector.
# The user-defined Python class implementing the Jacobi method.
class myJacobi:
# Setup the internal data. In this case, we access the matrix diagonal.
def setUp(self, pc):
_, P = pc.getOperators()
self.D = P.getDiagonal()
# Apply the preconditioner
def apply(self, pc, x, y):
y.pointwiseDivide(x, self.D)
We can run the script used to test our matrix class and use command line arguments to specify that our preconditioner should be used:
$ python mat.py -pc_type python -pc_python_type pc.myJacobi -ksp_view
KSP Object: 1 MPI process
type: cg
maximum iterations=10000, initial guess is zero
tolerances: relative=1e-05, absolute=1e-50, divergence=10000.
left preconditioning
using PRECONDITIONED norm type for convergence test
PC Object: 1 MPI process
type: python
Python: pc.myJacobi
linear system matrix = precond matrix:
Mat Object: 1 MPI process
type: python
rows=256, cols=256
Python: __main__.Poisson2D
PETSc Python linear solver type#
The protocol for the petsc4py.PETSc.KSP.Type.PYTHON Krylov solver is:
from petsc4py.PETSc import KSP
from petsc4py.PETSc import Vec
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by KSPPYTHON
class KSPPythonProtocol:
def solve(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Solve the linear system with right-hand side b. Return solution in x."""
...
def solveTranspose(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Solve the transposed linear system with right-hand side b. Return solution in x."""
...
def view(self, ksp: KSP, viewer: Viewer) -> None:
"""View the Krylov solver."""
...
def setFromOptions(self, ksp: KSP) -> None:
"""Process command line for customization."""
...
def setUp(self, ksp: KSP) -> None:
"""Perform the required setup."""
...
def buildSolution(self, ksp: KSP, x: Vec) -> None:
"""Compute the solution vector."""
...
def buildResidual(self, ksp: KSP, t: Vec, r: Vec) -> None:
"""Compute the residual vector, return it in r. t is a scratch working vector."""
...
def reset(self, ksp: KSP) -> None:
"""Reset the Krylov solver."""
...