Actual source code: ex83f.F90

  1: !
  2: !     Creates a tridiagonal sparse matrix explicitly in Fortran and solves a linear system with it
  3: !
  4: !     The matrix is provided in two three ways
  5: !       Compressed sparse row: ia(), ja(), and a()
  6: !     Entry triples:  rows(), cols(), and a()
  7: !     Entry triples in a way that supports new nonzero values with the same nonzero structure
  8: !
  9:       program main
 10: #include <petsc/finclude/petscksp.h>
 11:       use petscksp
 12:       implicit none

 14:       PetscInt i,n,nz
 15:       PetscBool flg,equal
 16:       PetscErrorCode ierr
 17:       PetscInt,ALLOCATABLE :: ia(:)
 18:       PetscInt,ALLOCATABLE :: ja(:)
 19:       PetscScalar,ALLOCATABLE :: a(:)
 20:       PetscScalar,ALLOCATABLE :: x(:)
 21:       PetscScalar,ALLOCATABLE :: b(:)

 23:       PetscInt,ALLOCATABLE :: rows(:)
 24:       PetscInt,ALLOCATABLE :: cols(:)

 26:       Mat J,Jt,Jr
 27:       Vec rhs,solution
 28:       KSP ksp
 29:       PC pc

 31:       PetscCallA(PetscInitialize(ierr))

 33:       n = 3
 34:       PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,'-n',n,flg,ierr))
 35:       nz = 3*n - 4;

 37:       ALLOCATE (b(n),x(n))

 39: !     Fill the sparse matrix representation
 40:       ALLOCATE (ia(n+1),ja(nz),a(nz))
 41:       ALLOCATE (rows(nz),cols(nz))

 43:       do i=1,n
 44:         b(i) = 1.0
 45:       enddo

 47: !     PETSc ia() and ja() values begin at 0, not 1, you may need to shift the indices used in your code
 48:       ia(1) = 0
 49:       ia(2) = 1
 50:       do i=3,n
 51:          ia(i) = ia(i-1) + 3
 52:       enddo
 53:       ia(n+1) = ia(n) + 1

 55:       ja(1) = 0
 56:       rows(1) = 0; cols(1) = 0
 57:       a(1)  = 1.0
 58:       do i=2,n-1
 59:          ja(2+3*(i-2))   = i-2
 60:          rows(2+3*(i-2)) = i-1; cols(2+3*(i-2)) = i-2
 61:          a(2+3*(i-2))    = -1.0;
 62:          ja(2+3*(i-2)+1) = i-1
 63:          rows(2+3*(i-2)+1) = i-1; cols(2+3*(i-2)+1) = i-1
 64:          a(2+3*(i-2)+1)  = 2.0;
 65:          ja(2+3*(i-2)+2) = i
 66:          rows(2+3*(i-2)+2) = i-1; cols(2+3*(i-2)+2) = i
 67:          a(2+3*(i-2)+2)  = -1.0;
 68:       enddo
 69:       ja(nz) = n-1
 70:       rows(nz) = n-1; cols(nz) = n-1
 71:       a(nz) = 1.0

 73:       PetscCallA(MatCreateSeqAIJWithArrays(PETSC_COMM_SELF,n,n,ia,ja,a,J,ierr))
 74:       PetscCallA(MatCreateSeqAIJFromTriple(PETSC_COMM_SELF,n,n,rows,cols,a,Jt,nz,PETSC_FALSE,ierr))
 75:       PetscCallA(MatEqual(J,Jt,equal,ierr))
 76:       if (equal .neqv. PETSC_TRUE) then
 77:          SETERRA(PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jt should be equal')
 78:       endif
 79:       PetscCallA(MatDestroy(Jt,ierr))
 80:       PetscCallA(MatCreate(PETSC_COMM_SELF,Jr,ierr))
 81:       PetscCallA(MatSetSizes(Jr,n,n,n,n,ierr))
 82:       PetscCallA(MatSetType(Jr,MATSEQAIJ,ierr))
 83:       PetscCallA(MatSetPreallocationCOO(Jr,nz,rows,cols,ierr))
 84:       PetscCallA(MatSetValuesCOO(Jr,a,INSERT_VALUES,ierr))
 85:       PetscCallA(MatEqual(J,Jr,equal,ierr))
 86:       if (equal .neqv. PETSC_TRUE) then
 87:          SETERRA(PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jr should be equal')
 88:       endif

 90:       PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF,1,n,b,rhs,ierr))
 91:       PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF,1,n,x,solution,ierr))

 93:       PetscCallA(KSPCreate(PETSC_COMM_SELF,ksp,ierr))
 94:       PetscCallA(KSPSetErrorIfNotConverged(ksp,PETSC_TRUE,ierr))
 95: !     Default to a direct sparse LU solver for robustness
 96:       PetscCallA(KSPGetPC(ksp,pc,ierr))
 97:       PetscCallA(PCSetType(pc,PCLU,ierr))
 98:       PetscCallA(KSPSetFromOptions(ksp,ierr))
 99:       PetscCallA(KSPSetOperators(ksp,J,J,ierr))

101:       PetscCallA(KSPSolve(ksp,rhs,solution,ierr))

103: !     Keep the same size and nonzero structure of the matrix but change its numerical entries
104:       do i=2,n-1
105:          a(2+3*(i-2)+1)  = 4.0;
106:       enddo
107:       PetscCallA(PetscObjectStateIncrease(J,ierr))
108:       PetscCallA(MatSetValuesCOO(Jr,a,INSERT_VALUES,ierr))
109:       PetscCallA(MatEqual(J,Jr,equal,ierr))
110:       if (equal .neqv. PETSC_TRUE) then
111:          SETERRA(PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jr should be equal')
112:       endif
113:       PetscCallA(MatDestroy(Jr,ierr))

115:       PetscCallA(KSPSolve(ksp,rhs,solution,ierr))

117:       PetscCallA(KSPDestroy(ksp,ierr))
118:       PetscCallA(VecDestroy(rhs,ierr))
119:       PetscCallA(VecDestroy(solution,ierr))
120:       PetscCallA(MatDestroy(J,ierr))

122:       DEALLOCATE (b,x)
123:       DEALLOCATE (ia,ja,a)
124:       DEALLOCATE (rows,cols)

126:       PetscCallA(PetscFinalize(ierr))
127:       end

129: !/*TEST
130: !
131: !     test:
132: !       args: -ksp_monitor -ksp_view
133: !
134: !TEST*/