1# Poisson in 2D 2# ============= 3# 4# Solve a constant coefficient Poisson problem on a regular grid. The source 5# code for this demo can be `downloaded here <../../_static/poisson2d.py>`__ 6# 7# .. math:: 8# 9# - u_{xx} - u_{yy} = 1 \quad\textsf{in}\quad [0,1]^2\\ 10# u = 0 \quad\textsf{on the boundary.} 11 12# This is a naïve, parallel implementation, using :math:`n` interior grid 13# points per dimension and a lexicographic ordering of the nodes. 14 15# This code is kept as simple as possible. However, simplicity comes at a 16# price. Here we use a naive decomposition that does not lead to an optimal 17# communication complexity for the matrix-vector product. An optimal complexity 18# decomposition of a structured grid could be achieved using `PETSc.DMDA`. 19# 20# This demo is structured as a script to be executed using: 21# 22# .. code-block:: console 23# 24# $ python poisson2d.py 25# 26# potentially with additional options passed at the end of the command. 27# 28# At the start of your script, call `petsc4py.init` passing `sys.argv` so that 29# command-line arguments to the script are passed through to PETSc. 30 31import sys 32import petsc4py 33 34petsc4py.init(sys.argv) 35 36# The full PETSc4py API is to be found in the `petsc4py.PETSc` module. 37 38from petsc4py import PETSc 39 40# PETSc is extensively programmable using the `PETSc.Options` database. For 41# more information see `working with PETSc Options <petsc_options>`. 42 43OptDB = PETSc.Options() 44 45# Grid size and spacing using a default value of ``5``. The user can specify a 46# different number of points in each direction by passing the ``-n`` option to 47# the script. 48 49n = OptDB.getInt('n', 5) 50h = 1.0 / (n + 1) 51 52# Matrices are instances of the `PETSc.Mat` class. 53 54A = PETSc.Mat() 55 56# Create the underlying PETSc C Mat object. 57# You can omit the ``comm`` argument if your objects live on 58# `PETSc.COMM_WORLD` but it is a dangerous choice to rely on default values 59# for such important arguments. 60 61A.create(comm=PETSc.COMM_WORLD) 62 63# Specify global matrix shape with a tuple. 64 65A.setSizes((n * n, n * n)) 66 67# The call above implicitly assumes that we leave the parallel decomposition of 68# the matrix rows to PETSc by using `PETSc.DECIDE` for local sizes. 69# It is equivalent to: 70# 71# .. code-block:: python 72# 73# A.setSizes(((PETSc.DECIDE, n * n), (PETSc.DECIDE, n * n))) 74 75# Here we use a sparse matrix of AIJ type 76# Various `matrix formats <petsc4py.PETSc.Mat.Type>` can be selected: 77 78A.setType(PETSc.Mat.Type.AIJ) 79 80# Finally we allow the user to set any options they want to on the matrix from 81# the command line: 82 83A.setFromOptions() 84 85# Insertion into some matrix types is vastly more efficient if we preallocate 86# space rather than allow this to happen dynamically. Here we hint the number 87# of nonzeros to be expected on each row. 88 89A.setPreallocationNNZ(5) 90 91# We can now write out our finite difference matrix assembly using conventional 92# Python syntax. `Mat.getOwnershipRange` is used to retrieve the range of rows 93# local to this processor. 94 95 96def index_to_grid(r): 97 """Convert a row number into a grid point.""" 98 return (r // n, r % n) 99 100 101rstart, rend = A.getOwnershipRange() 102for row in range(rstart, rend): 103 i, j = index_to_grid(row) 104 A[row, row] = 4.0 / h**2 105 if i > 0: 106 column = row - n 107 A[row, column] = -1.0 / h**2 108 if i < n - 1: 109 column = row + n 110 A[row, column] = -1.0 / h**2 111 if j > 0: 112 column = row - 1 113 A[row, column] = -1.0 / h**2 114 if j < n - 1: 115 column = row + 1 116 A[row, column] = -1.0 / h**2 117 118# At this stage, any exchange of information required in the matrix assembly 119# process has not occurred. We achieve this by calling `Mat.assemblyBegin` and 120# then `Mat.assemblyEnd`. 121 122A.assemblyBegin() 123A.assemblyEnd() 124 125# We set up an additional option so that the user can print the matrix by 126# passing ``-view_mat`` to the script. 127 128A.viewFromOptions('-view_mat') 129 130# PETSc represents all linear solvers as preconditioned Krylov subspace methods 131# of type `PETSc.KSP`. Here we create a KSP object for a conjugate gradient 132# solver preconditioned with an algebraic multigrid method. 133 134ksp = PETSc.KSP() 135ksp.create(comm=A.getComm()) 136ksp.setType(PETSc.KSP.Type.CG) 137ksp.getPC().setType(PETSc.PC.Type.GAMG) 138 139# We set the matrix in our linear solver and allow the user to program the 140# solver with options. 141 142ksp.setOperators(A) 143ksp.setFromOptions() 144 145# Since the matrix knows its size and parallel distribution, we can retrieve 146# appropriately-scaled vectors using `Mat.createVecs`. PETSc vectors are 147# objects of type `PETSc.Vec`. Here we set the right-hand side of our system to 148# a vector of ones, and then solve. 149 150x, b = A.createVecs() 151b.set(1.0) 152ksp.solve(b, x) 153 154# Finally, allow the user to print the solution by passing ``-view_sol`` to the 155# script. 156 157x.viewFromOptions('-view_sol') 158 159# Things to try 160# ------------- 161# 162# - Show the solution with ``-view_sol``. 163# - Show the matrix with ``-view_mat``. 164# - Change the resolution with ``-n``. 165# - Use a direct solver by passing ``-ksp_type preonly -pc_type lu``. 166# - Run in parallel on two processors using: 167# 168# .. code-block:: console 169# 170# mpiexec -n 2 python poisson2d.py 171