Poisson in 2D#
Solve a constant coefficient Poisson problem on a regular grid. The source code for this demo can be downloaded here
This is a naïve, parallel implementation, using \(n\) interior grid points per dimension and a lexicographic ordering of the nodes.
This code is kept as simple as possible. However, simplicity comes at a
price. Here we use a naive decomposition that does not lead to an optimal
communication complexity for the matrix-vector product. An optimal complexity
decomposition of a structured grid could be achieved using
This demo is structured as a script to be executed using:
$ python poisson2d.py
potentially with additional options passed at the end of the command.
import sys import petsc4py petsc4py.init(sys.argv)
The full PETSc4py API is to be found in the
from petsc4py import PETSc
OptDB = PETSc.Options()
Grid size and spacing using a default value of
5. The user can specify a
different number of points in each direction by passing the
-n option to
n = OptDB.getInt('n', 5) h = 1.0 / (n + 1)
Sparse matrices are represented by
You can omit the
comm argument if your objects live on
PETSc.COMM_WORLD but it is a dangerous choice to rely on default values
for such important arguments.
A = PETSc.Mat() A.create(comm=PETSc.COMM_WORLD)
Specify global matrix shape with a tuple.
A.setSizes((n * n, n * n))
The call above implicitly assumes that we leave the parallel decomposition of
the matrix rows to PETSc by using
PETSc.DECIDE for local sizes.
It is equivalent to:
A.setSizes(((PETSc.DECIDE, n * n), (PETSc.DECIDE, n * n)))
sparse matrix formats can be selected:
Finally we allow the user to set any options they want to on the matrix from the command line:
Insertion into some matrix types is vastly more efficient if we preallocate space rather than allow this to happen dynamically. Here we hint the number of nonzeros to be expected on each row.
We can now write out our finite difference matrix assembly using conventional
Mat.getOwnershipRange is used to retrieve the range of rows
local to this processor.
def index_to_grid(r): """Convert a row number into a grid point.""" return (r // n, r % n) rstart, rend = A.getOwnershipRange() for row in range(rstart, rend): i, j = index_to_grid(row) A[row, row] = 4.0 / h**2 if i > 0: column = row - n A[row, column] = -1.0 / h**2 if i < n - 1: column = row + n A[row, column] = -1.0 / h**2 if j > 0: column = row - 1 A[row, column] = -1.0 / h**2 if j < n - 1: column = row + 1 A[row, column] = -1.0 / h**2
We set up an additional option so that the user can print the matrix by
-view_mat to the script.
PETSc represents all linear solvers as preconditioned Krylov subspace methods
PETSc.KSP. Here we create a KSP object for a conjugate gradient
solver preconditioned with an algebraic multigrid method.
ksp = PETSc.KSP() ksp.create(comm=A.getComm()) ksp.setType(PETSc.KSP.Type.CG) ksp.getPC().setType(PETSc.PC.Type.GAMG)
We set the matrix in our linear solver and allow the user to program the solver with options.
Since the matrix knows its size and parallel distribution, we can retrieve
appropriately-scaled vectors using
Mat.createVecs. PETSc vectors are
objects of type
PETSc.Vec. Here we set the right hand side of our system to
a vector of ones, and then solve.
x, b = A.createVecs() b.set(1.0) ksp.solve(b, x)
Finally, allow the user to print the solution by passing
-view_sol to the
Things to try#
Show the solution with
Show the matrix with
Change the resolution with
Use a direct solver by passing
-ksp_type preonly -pc_type lu.
Run in parallel on two processors using:
mpiexec -n 2 python poisson2d.py