# 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 contained in

`\$PETSC_DIR/src/binding/petsc4py/src/petsc4py/PETSc/libpetsc4py.pyx`

The scripts can be found in `\$PETSC_DIR/src/binding/petsc4py/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: sMat) -> 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, 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):
assert da.getDim() == 2
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 Mat
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."""
...

```