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: #include <petsc/finclude/petscksp.h>
 10: program main
 11:   use petscksp
 12:   implicit none

 14:   PetscInt i, n
 15:   PetscCount nz
 16:   PetscBool flg, equal
 17:   PetscErrorCode ierr
 18:   PetscInt, allocatable, dimension(:) :: ia, ja, rows, cols
 19:   PetscScalar, allocatable, dimension(:) :: a, x, b

 21:   Mat J, Jt, Jr
 22:   Vec rhs, solution
 23:   KSP ksp
 24:   PC pc

 26:   PetscCallA(PetscInitialize(ierr))
 27:   n = 3
 28:   PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-n', n, flg, ierr))
 29:   nz = 3*n - 4

 31:   allocate (b(n), x(n))

 33: !     Fill the sparse matrix representation
 34:   allocate (ia(n + 1), ja(nz), a(nz))
 35:   allocate (rows(nz), cols(nz))

 37:   b(1:n) = 1.0

 39: !     PETSc ia() and ja() values begin at 0, not 1, you may need to shift the indices used in your code
 40:   ia(1) = 0
 41:   ia(2) = 1
 42:   do i = 3, n
 43:     ia(i) = ia(i - 1) + 3
 44:   end do
 45:   ia(n + 1) = ia(n) + 1

 47:   ja(1) = 0
 48:   rows(1) = 0; cols(1) = 0
 49:   a(1) = 1.0
 50:   do i = 2, n - 1
 51:     ja(2 + 3*(i - 2)) = i - 2
 52:     rows(2 + 3*(i - 2)) = i - 1; cols(2 + 3*(i - 2)) = i - 2
 53:     a(2 + 3*(i - 2)) = -1.0
 54:     ja(2 + 3*(i - 2) + 1) = i - 1
 55:     rows(2 + 3*(i - 2) + 1) = i - 1; cols(2 + 3*(i - 2) + 1) = i - 1
 56:     a(2 + 3*(i - 2) + 1) = 2.0
 57:     ja(2 + 3*(i - 2) + 2) = i
 58:     rows(2 + 3*(i - 2) + 2) = i - 1; cols(2 + 3*(i - 2) + 2) = i
 59:     a(2 + 3*(i - 2) + 2) = -1.0
 60:   end do
 61:   ja(nz) = n - 1
 62:   rows(nz) = n - 1; cols(nz) = n - 1
 63:   a(nz) = 1.0

 65:   PetscCallA(MatCreateSeqAIJWithArrays(PETSC_COMM_SELF, n, n, ia, ja, a, J, ierr))
 66:   PetscCallA(MatCreateSeqAIJFromTriple(PETSC_COMM_SELF, n, n, rows, cols, a, Jt, nz, PETSC_FALSE, ierr))
 67:   PetscCallA(MatEqual(J, Jt, equal, ierr))
 68:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jt must be equal')
 69:   PetscCallA(MatDestroy(Jt, ierr))
 70:   PetscCallA(MatCreate(PETSC_COMM_SELF, Jr, ierr))
 71:   PetscCallA(MatSetSizes(Jr, n, n, n, n, ierr))
 72:   PetscCallA(MatSetType(Jr, MATSEQAIJ, ierr))
 73:   PetscCallA(MatSetPreallocationCOO(Jr, nz, rows, cols, ierr))
 74:   PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
 75:   PetscCallA(MatEqual(J, Jr, equal, ierr))
 76:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr must be equal')

 78:   PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, 1_PETSC_INT_KIND, n, b, rhs, ierr))
 79:   PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, 1_PETSC_INT_KIND, n, x, solution, ierr))

 81:   PetscCallA(KSPCreate(PETSC_COMM_SELF, ksp, ierr))
 82:   PetscCallA(KSPSetErrorIfNotConverged(ksp, PETSC_TRUE, ierr))
 83: ! Default to a direct sparse LU solver for robustness
 84:   PetscCallA(KSPGetPC(ksp, pc, ierr))
 85:   PetscCallA(PCSetType(pc, PCLU, ierr))
 86:   PetscCallA(KSPSetFromOptions(ksp, ierr))
 87:   PetscCallA(KSPSetOperators(ksp, J, J, ierr))

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

 91: ! Keep the same size and nonzero structure of the matrix but change its numerical entries
 92:   do i = 2, n - 1
 93:     a(2 + 3*(i - 2) + 1) = 4.0
 94:   end do
 95:   PetscCallA(PetscObjectStateIncrease(J, ierr))
 96:   PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
 97:   PetscCallA(MatEqual(J, Jr, equal, ierr))
 98:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr should be equal')
 99:   PetscCallA(MatDestroy(Jr, ierr))

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

103:   PetscCallA(KSPDestroy(ksp, ierr))
104:   PetscCallA(VecDestroy(rhs, ierr))
105:   PetscCallA(VecDestroy(solution, ierr))
106:   PetscCallA(MatDestroy(J, ierr))

108:   deallocate (b, x)
109:   deallocate (ia, ja, a)
110:   deallocate (rows, cols)

112:   PetscCallA(PetscFinalize(ierr))
113: end

115: !/*TEST
116: !
117: !     test:
118: !       args: -ksp_monitor -ksp_view
119: !
120: !TEST*/