1 2''' 3Ex23 from PETSc example files implemented for PETSc4py. 4https://petsc.org/release/src/ksp/ksp/tutorials/ex23.c.html 5By: Miguel Arriaga 6 7Solves a tridiagonal linear system. 8 9Vec x, b, u; approx solution, RHS, exact solution 10Mat A; linear system matrix 11KSP ksp; linear solver context 12PC pc; preconditioner context 13PetscReal norm; norm of solution error 14 15''' 16import sys 17import petsc4py 18petsc4py.init(sys.argv) 19from petsc4py import PETSc 20 21import numpy as np 22 23comm = PETSc.COMM_WORLD 24size = comm.getSize() 25rank = comm.getRank() 26n = 12 # Size of problem 27tol = 1E-11 # Tolerance of Result. tol=1000.*PETSC_MACHINE_EPSILON; 28 29''' 30- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 31 Compute the matrix and right-hand-side vector that define 32 the linear system, Ax = b. 33- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 34 35Create vectors. Note that we form 1 vector from scratch and 36then duplicate as needed. For this simple case let PETSc decide how 37many elements of the vector are stored on each processor. The second 38argument to VecSetSizes() below causes PETSc to decide. 39''' 40 41x = PETSc.Vec().create(comm=comm) 42x.setSizes(n) 43x.setFromOptions() 44 45b = x.duplicate() 46u = x.duplicate() 47 48''' 49Identify the starting and ending mesh points on each 50processor for the interior part of the mesh. We let PETSc decide 51above. 52''' 53 54rstart,rend = x.getOwnershipRange() 55nlocal = x.getLocalSize() 56 57''' 58Create matrix. When using MatCreate(), the matrix format can 59be specified at runtime. 60 61Performance tuning note: For problems of substantial size, 62preallocation of matrix memory is crucial for attaining good 63performance. See the matrix chapter of the users manual for details. 64 65We pass in nlocal as the "local" size of the matrix to force it 66to have the same parallel layout as the vector created above. 67''' 68 69A = PETSc.Mat().create(comm=comm) 70A.setSizes(n,nlocal) 71A.setFromOptions() 72A.setUp() 73 74''' 75Assemble matrix. 76 77The linear system is distributed across the processors by 78chunks of contiguous rows, which correspond to contiguous 79sections of the mesh on which the problem is discretized. 80For matrix assembly, each processor contributes entries for 81the part that it owns locally. 82''' 83 84col = np.zeros(3,dtype=PETSc.IntType) 85value = np.zeros(3,dtype=PETSc.ScalarType) 86 87if not rstart: 88 rstart = 1 89 i = 0; col[0] = 0; col[1] = 1; value[0] = 2.0; value[1] = -1.0 90 A.setValues(i,col[0:2],value[0:2]) 91 92if rend == n: 93 rend = n-1 94 i = n-1; col[0] = n-2; col[1] = n-1; value[0] = -1.0; value[1] = 2.0 95 A.setValues(i,col[0:2],value[0:2]) 96 97 98''' Set entries corresponding to the mesh interior ''' 99value[0] = -1.0; value[1] = 2.0; value[2] = -1.0 100for i in range(rstart,rend): 101 col[0] = i-1; col[1] = i; col[2] = i+1 102 A.setValues(i,col,value) 103 104 105''' Assemble the matrix ''' 106A.assemblyBegin(A.AssemblyType.FINAL) 107A.assemblyEnd(A.AssemblyType.FINAL) 108 109''' Set exact solution; then compute right-hand-side vector. ''' 110u.set(1.0) 111b = A(u) 112 113''' 114- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 115 Create the linear solver and set various options 116- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - ''' 117''' 118Create linear solver context 119''' 120ksp = PETSc.KSP().create() 121 122''' 123Set operators. Here the matrix that defines the linear system 124also serves as the matrix from which the preconditioner is constructed. 125''' 126ksp.setOperators(A,A) 127 128''' 129Set linear solver defaults for this problem (optional). 130 - By extracting the KSP and PC contexts from the KSP context, 131 we can then directly call any KSP and PC routines to set 132 various options. 133 - The following four statements are optional; all of these 134 parameters could alternatively be specified at runtime via 135 KSPSetFromOptions(); 136''' 137pc = ksp.getPC() 138pc.setType('jacobi') 139ksp.setTolerances(rtol=1.e-7) 140 141''' 142Set runtime options, e.g., 143-ksp_type <type> -pc_type <type> -ksp_monitor -ksp_rtol <rtol> 144These options will override those specified above as long as 145KSPSetFromOptions() is called _after_ any other customization 146routines. 147''' 148ksp.setFromOptions() 149 150''' 151- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 152 Solve the linear system 153- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - ''' 154 155''' Solve linear system ''' 156ksp.solve(b,x) 157 158''' 159View solver info; we could instead use the option -ksp_view to 160print this info to the screen at the conclusion of KSPSolve(). 161''' 162# ksp.view() 163 164''' 165- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 166 Check solution and clean up 167- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - ''' 168 169''' Check the error ''' 170x = x - u # x.axpy(-1.0,u) 171norm = x.norm(PETSc.NormType.NORM_2) 172its = ksp.getIterationNumber() 173if norm > tol: 174 PETSc.Sys.Print("Norm of error {}, Iterations {}\n".format(norm,its),comm=comm) 175else: 176 if size==1: 177 PETSc.Sys.Print("- Serial OK",comm=comm) 178 else: 179 PETSc.Sys.Print("- Parallel OK",comm=comm) 180