# 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."""
            ...

## PETSc Python nonlinear solver type (TODO)

## PETSc Python ode-integrator type (TODO)

## PETSc Python optimization solver type (TODO)