PETSc Python types#
PETSc supports Python implementations of matrices, preconditioners, Krylov solvers, nonlinear solvers, ODE integrators, optimizers, and viewers.
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.
Each implementation is a Python object, called a context, whose methods PETSc calls to perform the supported operations. The protocol classes below describe the callback signatures; they are not base classes to inherit from. Implement only the callbacks your context needs, and omit unused methods. An empty method still counts as an implementation and can override PETSc’s default behavior.
The optional create(obj) callback runs when the context is attached to a
PETSc object. The optional destroy(obj) callback releases resources before
the context is replaced or removed, including when the PETSc object is
destroyed. These callbacks are distinct from setUp(obj) and, where
supported, reset(obj), which prepare and reset an object for reuse.
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.typing import ArrayInt
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 create(self, A: Mat) -> None:
"""Initialize resources when the context is attached to the matrix."""
...
def destroy(self, A: Mat) -> None:
"""Release resources when the context is detached from the matrix."""
...
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 options from the options database."""
...
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 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:
"""Add s to each diagonal entry of the matrix."""
...
def createSubMatrix(self, A: Mat, r: IS, c: IS, out: Mat | None) -> Mat:
"""Return the submatrix corresponding to r rows and c columns.
Matrix out must be reused if not None.
"""
...
def zeroRows(self, A: Mat, r: ArrayInt, diag: Scalar, x: Vec, b: Vec) -> None:
"""Zero the rows indexed by the NumPy integer array r.
The array owns a copy of this process's global row indices and may
be retained after the callback returns.
Insert diag on the diagonal. If x and b are provided, use the known
solution values in x to adjust b. Missing vectors are empty Vec
wrappers, not None; check them with bool(x) and bool(b).
"""
...
def zeroRowsColumns(
self, A: Mat, r: ArrayInt, diag: Scalar, x: Vec, b: Vec
) -> None:
"""Zero the rows and columns indexed by the NumPy integer array r.
The array owns a copy of this process's global row indices and may
be retained after the callback returns.
Insert diag on the diagonal. If x and b are provided, use the known
solution values in x to adjust b. Missing vectors are empty Vec
wrappers, not None; check them with bool(x) and bool(b).
"""
...
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 diagonalScale(self, A: Mat, L: Vec, R: Vec) -> None:
"""Scale the matrix on the left by L and on the right by R.
A = diag(L) @ A @ diag(R). A missing vector is an empty Vec wrapper,
not None. Omit scaling on a side when its vector evaluates to False.
"""
...
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 | None
) -> bool:
"""Return whether the matrix supports the requested product type.
Z is None for all product types except ABC.
"""
...
def productSymbolic(
self, A: Mat, product: Mat, producttype: str, X: Mat, Y: Mat, Z: Mat | None
) -> None:
"""Perform the symbolic stage of the requested matrix product.
Z is None for all product types except ABC.
"""
...
def productNumeric(
self, A: Mat, product: Mat, producttype: str, X: Mat, Y: Mat, Z: Mat | None
) -> None:
"""Perform the numeric stage of the requested matrix product.
Z is None for all product types except ABC.
"""
...
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:
"""Replace each entry by its imaginary part: A = imag(A)."""
...
def realPart(self, A: Mat) -> None:
"""Replace each entry by its real part: 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 < my - 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 create(self, pc: PC) -> None:
"""Initialize resources when the context is attached to the PC."""
...
def destroy(self, pc: PC) -> None:
"""Release resources when the context is detached from the PC."""
...
def apply(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the preconditioner to b, storing the result in x."""
...
def applySymmetricLeft(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the symmetric left part to b, storing the result in x."""
...
def applySymmetricRight(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the symmetric right part to b, storing the result in x."""
...
def applyTranspose(self, pc: PC, b: Vec, x: Vec) -> None:
"""Apply the transpose to b, storing the result in x."""
...
def matApply(self, pc: PC, B: Mat, X: Mat) -> None:
"""Apply the preconditioner to B, storing the result in X."""
...
def preSolve(self, pc: PC, ksp: KSP, b: Vec, x: Vec) -> None:
"""Prepare for a Krylov solve.
This method may modify the right-hand side b and initial guess x.
"""
...
def postSolve(self, pc: PC, ksp: KSP, b: Vec, x: Vec) -> None:
"""Postprocess a Krylov solve.
This method may modify the right-hand side b and solution x.
"""
...
def view(self, pc: PC, viewer: Viewer) -> None:
"""View the preconditioner."""
...
def setFromOptions(self, pc: PC) -> None:
"""Process options from the options database."""
...
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)
From demo/python_types, 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, which is also used to construct the preconditioner:
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 create(self, ksp: KSP) -> None:
"""Initialize resources when the context is attached to the KSP."""
...
def destroy(self, ksp: KSP) -> None:
"""Release resources when the context is detached from the KSP."""
...
def solve(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Solve with right-hand side b, storing the solution in x.
Implement this method to control the complete solve and set its
convergence reason. Omit it to use the default iteration with step(),
preStep(), and postStep().
"""
...
def solveTranspose(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Solve the transposed system with right-hand side b and solution x.
Implement this method to control the complete solve and set its
convergence reason. Omit it to use the default iteration with
stepTranspose(), preStep(), and postStep().
"""
...
def step(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Update x for right-hand side b in the default iteration."""
...
def stepTranspose(self, ksp: KSP, b: Vec, x: Vec) -> None:
"""Update x for the transposed system in the default iteration."""
...
def preStep(self, ksp: KSP) -> None:
"""Process the solver state before a default iteration."""
...
def postStep(self, ksp: KSP) -> None:
"""Process the solver state after a default iteration."""
...
def view(self, ksp: KSP, viewer: Viewer) -> None:
"""View the Krylov solver."""
...
def setFromOptions(self, ksp: KSP) -> None:
"""Process options from the options database."""
...
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 in r, using t as a work vector."""
...
def reset(self, ksp: KSP) -> None:
"""Reset the Krylov solver."""
...
The following example implements one step of preconditioned Richardson iteration:
\[r_k = b - A x_k, \qquad z_k = P^{-1} r_k, \qquad
x_{k+1} = x_k + \omega z_k.\]
The step() method updates the solution. By omitting solve(), the context
uses the default KSPPYTHON loop to compute residuals, check convergence, and
call monitors. This example uses left preconditioning and the unpreconditioned
residual norm. Work vectors are allocated in setUp() and released in
reset() and destroy(). The relaxation factor defaults to \(\omega = 1\)
and can be changed with -ksp_richardson_scale omega as implemented in
setFromOptions().
from petsc4py import PETSc
# The user-defined Python class implementing preconditioned Richardson iteration.
class Richardson:
def __init__(self):
self.omega = 1.0
self.work = []
def create(self, ksp):
# The default loop measures the norm of b - A x.
ksp.setPCSide(PETSc.PC.Side.LEFT)
ksp.setNormType(PETSc.KSP.NormType.UNPRECONDITIONED)
def setFromOptions(self, ksp):
options = PETSc.Options(ksp.getOptionsPrefix())
self.omega = options.getReal('ksp_richardson_scale', self.omega)
def view(self, ksp, viewer):
if viewer.getType() == PETSc.Viewer.Type.ASCII:
viewer.printfASCII(f' relaxation factor: {self.omega:g}\n')
def setUp(self, ksp):
self.reset(ksp)
self.work = ksp.getWorkVecs(right=2)
def reset(self, ksp):
for vec in self.work:
vec.destroy()
self.work = []
def destroy(self, ksp):
self.reset(ksp)
def step(self, ksp, b, x):
A, _ = ksp.getOperators()
z, r = self.work
A.mult(x, r)
r.aypx(-1, b)
ksp.getPC().apply(r, z)
x.axpy(self.omega, z)
From demo/python_types, we can run the matrix example with the Python
Richardson solver and Jacobi preconditioning. The -ksp_view option displays
the solver type, Python context, and relaxation factor reported by view():
$ python mat.py -ksp_type python -ksp_python_type ksp.Richardson \
-pc_type jacobi -ksp_view
KSP Object: 1 MPI process
type: python
Python: ksp.Richardson
relaxation factor: 1
maximum iterations=10000, initial guess is zero
tolerances: relative=1e-05, absolute=1e-50, divergence=10000.
left preconditioning
using UNPRECONDITIONED norm type for convergence test
PC Object: 1 MPI process
type: jacobi
type DIAGONAL
linear system matrix, which is also used to construct the preconditioner:
Mat Object: 1 MPI process
type: python
rows=256, cols=256
Python: __main__.Poisson2D
PETSc Python nonlinear solver type#
The protocol for the petsc4py.PETSc.SNES.Type.PYTHON nonlinear solver is:
from petsc4py.PETSc import SNES
from petsc4py.PETSc import Vec
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by SNESPYTHON
class SNESPythonProtocol:
def create(self, snes: SNES) -> None:
"""Initialize resources when the context is attached to the SNES."""
...
def destroy(self, snes: SNES) -> None:
"""Release resources when the context is detached from the SNES."""
...
def solve(self, snes: SNES, b: Vec | None, x: Vec) -> None:
"""Solve the nonlinear system with a user-defined routine.
Implement this method to control the complete solve and set its
convergence reason. Omit it to use step(), preStep(), and postStep()
to customize the default solve. Store the solution in x.
"""
...
def step(self, snes: SNES, x: Vec, f: Vec, y: Vec) -> None:
"""Compute update y from x and residual f in the default solve."""
...
def preStep(self, snes: SNES) -> None:
"""Process the solver state before each step in the default solve."""
...
def postStep(self, snes: SNES) -> None:
"""Process the solver state after each step in the default solve."""
...
def view(self, snes: SNES, viewer: Viewer) -> None:
"""View the nonlinear solver."""
...
def setFromOptions(self, snes: SNES) -> None:
"""Process options from the options database."""
...
def setUp(self, snes: SNES) -> None:
"""Perform the required setup."""
...
def reset(self, snes: SNES) -> None:
"""Reset the nonlinear solver."""
...
The following example implements the complete nonlinear solve with
scipy.optimize.root. It solves \(x^2 - 2 = 0\) on one process.
import sys
import numpy as np
import petsc4py
from scipy.optimize import root
petsc4py.init(sys.argv)
from petsc4py import PETSc
# The user-defined Python class implementing the nonlinear solve with SciPy.
class SciPyRoot:
# Solve the complete nonlinear system.
def solve(self, snes, b, x):
# Create PETSc vectors used to evaluate the user-provided function.
work_x = x.duplicate()
work_f = x.duplicate()
# Adapt the PETSc function callback to the interface expected by SciPy.
def function(values):
work_x.setArray(values)
snes.computeFunction(work_x, work_f)
return work_f.getArray(readonly=True)
# Report an inner-solver failure unless SciPy completes successfully.
reason = PETSc.SNES.ConvergedReason.DIVERGED_INNER
message = None
try:
# Solve the system and return the solution to PETSc.
result = root(function, x.getArray(readonly=True))
x[:] = result.x[:]
# Record the solver state in the PETSc SNES object.
snes.setIterationNumber(result.nfev)
snes.setFunctionNorm(float(np.linalg.norm(result.fun)))
if result.success:
reason = PETSc.SNES.ConvergedReason.CONVERGED_ITS
else:
message = str(result.message)
except Exception as error:
message = str(error) or type(error).__name__
finally:
# Clean up work vectors.
work_x.destroy()
work_f.destroy()
snes.setConvergedReason(reason)
# Handle -snes_error_if_not_converged here instead of relying on PETSc's
# standard nonconvergence error. Raising a Python exception lets PETSc preserve
# SciPy's failure message when it converts the exception to a PETSc error.
if reason < 0 and snes.getErrorIfNotConverged():
if message:
msg = f'SciPy solve failed: {message}'
raise RuntimeError(msg)
raise RuntimeError('SciPy solve failed')
# Define the nonlinear equation x**2 - 2 = 0.
def function(snes, x, f):
f[0] = x[0] ** 2 - 2.0
# Create the initial guess and residual vector.
x = PETSc.Vec().createSeq(1, comm=PETSc.COMM_SELF)
x[0] = 1.0
f = x.duplicate()
# Create the Python SNES, set the equation, and solve it.
snes = PETSc.SNES().createPython(SciPyRoot(), comm=PETSc.COMM_SELF)
snes.setFunction(function, f)
snes.setErrorIfNotConverged()
snes.setFromOptions()
snes.solve(None, x)
PETSc.Sys.Print(f'sqrt(2) = {x[0]:.12f}')
PETSc Python ODE integrator type#
The protocol for the petsc4py.PETSc.TS.Type.PYTHON ODE integrator is:
from petsc4py.PETSc import Mat
from petsc4py.PETSc import SNES
from petsc4py.PETSc import TS
from petsc4py.PETSc import Vec
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by TSPYTHON
class TSPythonProtocol:
def create(self, ts: TS) -> None:
"""Initialize resources when the context is attached to the TS."""
...
def destroy(self, ts: TS) -> None:
"""Release resources when the context is detached from the TS."""
...
def step(self, ts: TS) -> None:
"""Advance the solution with a user-defined time-stepping routine.
Implement this method to control the complete time step. Omit it to
use solveStep() and adaptStep() to customize the default routine.
"""
...
def rollback(self, ts: TS) -> None:
"""Roll back the time integrator's internal state by one step."""
...
def interpolate(self, ts: TS, t: float, x: Vec) -> None:
"""Interpolate the solution at time t. Return the solution in x."""
...
def evaluatestep(self, ts: TS, order: int, x: Vec) -> bool:
"""Evaluate the current step at a given order.
Return the solution in x and whether the evaluation was available.
"""
...
def formSNESFunction(self, snes: SNES, x: Vec, f: Vec, ts: TS) -> None:
"""Form the SNES residual f at the candidate solution x."""
...
def formSNESJacobian(self, snes: SNES, x: Vec, A: Mat, B: Mat, ts: TS) -> None:
"""Form the Jacobian matrices A and B for the SNES solve at x."""
...
def solveStep(self, ts: TS, t: float, x: Vec) -> None:
"""Solve for the candidate solution x at time t."""
...
def adaptStep(
self, ts: TS, t: float, x: Vec
) -> tuple[float, bool] | float | bool | None:
"""Choose the next time-step size and whether to accept x.
Return the next time-step size and a flag: True to accept x, or False
to reject it. A float accepts x and sets the next time-step size.
A bool accepts or rejects x without changing the time-step size.
None accepts x without changing the time-step size.
"""
...
def view(self, ts: TS, viewer: Viewer) -> None:
"""View the time integrator."""
...
def setFromOptions(self, ts: TS) -> None:
"""Process options from the options database."""
...
def setUp(self, ts: TS) -> None:
"""Perform the required setup."""
...
def reset(self, ts: TS) -> None:
"""Reset the time integrator."""
...
The following example implements each time step with
scipy.integrate.odeint. It solves
\[\frac{du}{dt} = -u, \qquad u(0) = 1\]
on one process.
import sys
import petsc4py
from scipy.integrate import odeint
petsc4py.init(sys.argv)
from petsc4py import PETSc
# The user-defined Python class implementing each time step with SciPy.
class SciPyODEInt:
def __init__(self):
self.solution = PETSc.Vec()
self.function = PETSc.Vec()
# Create PETSc vectors used to evaluate the user-provided ODE function.
def setUp(self, ts):
self.solution = ts.getSolution().duplicate()
self.function = ts.getSolution().duplicate()
# Advance the solution from the current time to t.
def solveStep(self, ts, t, x):
# Adapt the PETSc ODE callback to the interface expected by SciPy.
def rhs(values, time):
self.solution.setArray(values)
ts.computeRHSFunction(time, self.solution, self.function)
return self.function.getArray(readonly=True)
# Integrate over the current PETSc time step and return the result.
initial = ts.getSolution().getArray(readonly=True)
solution = odeint(rhs, initial, [ts.getTime(), t])
x[0] = solution[-1]
# Destroy the work vectors when PETSc resets the integrator.
def reset(self, ts):
self.solution.destroy()
self.function.destroy()
# Define the ODE du/dt = -u.
def rhs(ts, t, x, f):
x.copy(f)
f.scale(-1.0)
# Create the initial condition.
u = PETSc.Vec().createSeq(1, comm=PETSc.COMM_SELF)
u[0] = 1.0
# Create the Python TS and configure the time integration.
ts = PETSc.TS().createPython(SciPyODEInt(), comm=PETSc.COMM_SELF)
ts.setRHSFunction(rhs)
ts.setSolution(u)
ts.setTime(0.0)
ts.setTimeStep(0.125)
ts.setMaxTime(1.0)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
ts.setFromOptions()
ts.solve(u)
PETSc.Sys.Print(f'u({ts.getTime():.1f}) = {u[0]:.12f}')
PETSc Python optimization solver type#
The protocol for the petsc4py.PETSc.TAO.Type.PYTHON TAO optimizer is:
from petsc4py.PETSc import TAO
from petsc4py.PETSc import Vec
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by TAOPYTHON
class TAOPythonProtocol:
def create(self, tao: TAO) -> None:
"""Initialize resources when the context is attached to the TAO."""
...
def destroy(self, tao: TAO) -> None:
"""Release resources when the context is detached from the TAO."""
...
def setFromOptions(self, tao: TAO) -> None:
"""Process options from the options database."""
...
def setUp(self, tao: TAO) -> None:
"""Set up the optimizer."""
...
def solve(self, tao: TAO) -> None:
"""Solve the optimization problem with a user-defined routine.
Implement this method to control the complete solve and set its
convergence reason. Omit it to use step(), preStep(), and postStep()
to customize the default solve.
"""
...
def step(self, tao: TAO, x: Vec, g: Vec | None, s: Vec | None) -> None:
"""Compute gradient g and search direction s at x.
Both g and s are None when the optimizer has no gradient routine.
"""
...
def preStep(self, tao: TAO) -> None:
"""Process the optimizer state before a default step."""
...
def postStep(self, tao: TAO) -> None:
"""Process the optimizer state after a default step."""
...
def view(self, tao: TAO, viewer: Viewer) -> None:
"""View the optimizer."""
...
The following example implements a gradient descent solver to minimize \(f(x) = (x_0 - 1)^2 + (x_1 - 2)^2\) on one process, starting from \(x = (0.5, 0.5)\).
It uses a petsc4py.PETSc.TAOLineSearch.Type.UNIT line search with step size
\(0.2\), giving the update
\[x^{k+1} = x^k - 0.2 \nabla f(x^k).\]
import sys
import petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
# The user-defined Python class implementing gradient descent.
class myGradientDescent:
def create(self, tao):
# Create a line search type with constant step size.
self._ls = PETSc.TAOLineSearch().create(comm=PETSc.COMM_SELF)
self._ls.useTAORoutine(tao)
self._ls.setType(PETSc.TAOLineSearch.Type.UNIT)
self._ls.setInitialStepLength(0.2)
def destroy(self, tao):
self._ls.destroy()
def solve(self, tao):
# Evaluate the objective and gradient at the initial guess.
x = tao.getSolution()
gradient = tao.getGradient()[0]
f = tao.computeObjectiveGradient(x, gradient)
tao.monitor(f=f, res=gradient.norm(), step=1.0)
if tao.checkConverged() != PETSc.TAO.ConvergedReason.CONTINUE_ITERATING:
return
# Prepare search direction for line search.
search_direction = gradient.duplicate()
# Optimization loop.
for it in range(1, tao.getMaximumIterations() + 1):
# Search in the negative gradient direction.
gradient.copy(search_direction)
search_direction.scale(-1)
# Update x and evaluate the new objective and gradient.
f, step, reason = self._ls.apply(x, gradient, search_direction)
if reason < 0:
tao.setConvergedReason(PETSc.TAO.ConvergedReason.DIVERGED_LS_FAILURE)
break
# Report the completed step and check for convergence or divergence.
tao.setIterationNumber(it)
tao.monitor(f=f, res=gradient.norm(), step=step)
if tao.checkConverged() != PETSc.TAO.ConvergedReason.CONTINUE_ITERATING:
break
if tao.getConvergedReason() == PETSc.TAO.ConvergedReason.CONTINUE_ITERATING:
tao.setConvergedReason(PETSc.TAO.ConvergedReason.DIVERGED_MAXITS)
search_direction.destroy()
# Minimize f(x) = (x[0] - 1)^2 + (x[1] - 2)^2.
def objective(tao, x):
return (x[0] - 1.0) ** 2 + (x[1] - 2.0) ** 2
def gradient(tao, x, g):
g[0] = 2.0 * (x[0] - 1.0)
g[1] = 2.0 * (x[1] - 2.0)
g.assemble()
# Create the initial guess and gradient vector.
x = PETSc.Vec().createSeq(2, comm=PETSc.COMM_SELF)
x.set(0.5)
g = x.duplicate()
# Create the Python optimizer and configure the minimization.
tao = PETSc.TAO().createPython(myGradientDescent(), comm=PETSc.COMM_SELF)
tao.setObjective(objective)
tao.setGradient(gradient, g)
tao.setSolution(x)
tao.setTolerances(gatol=1e-6)
tao.setMaximumIterations(100)
tao.setFromOptions()
tao.solve()
PETSc.Sys.Print(f'x = ({x[0]:.6f}, {x[1]:.6f})')
PETSc Python viewer#
The protocol for the petsc4py.PETSc.Viewer.Type.PYTHON viewer is:
from petsc4py.PETSc import Object
from petsc4py.PETSc import Viewer
# A template class with the Python methods supported by PETSCVIEWERPYTHON
class PetscViewerPythonProtocol:
def create(self, viewer: Viewer) -> None:
"""Initialize resources when the context is attached to the viewer."""
...
def destroy(self, viewer: Viewer) -> None:
"""Release resources when the context is detached from the viewer."""
...
def viewObject(self, viewer: Viewer, obj: Object) -> None:
"""View a generic object."""
...
def setUp(self, viewer: Viewer) -> None:
"""Set up the viewer."""
...
def setFromOptions(self, viewer: Viewer) -> None:
"""Process options from the options database."""
...
def flush(self, viewer: Viewer) -> None:
"""Flush the viewer."""
...
def view(self, viewer: Viewer, outviewer: Viewer) -> None:
"""View the viewer."""
...