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, one
     15:   PetscCount nz
     16:   PetscBool flg, equal
     17:   PetscErrorCode ierr
     18:   PetscInt, ALLOCATABLE :: ia(:)
     19:   PetscInt, ALLOCATABLE :: ja(:)
     20:   PetscScalar, ALLOCATABLE :: a(:)
     21:   PetscScalar, ALLOCATABLE :: x(:)
     22:   PetscScalar, ALLOCATABLE :: b(:)

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

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

     32:   PetscCallA(PetscInitialize(ierr))
     33:   one = 1
     34:   n = 3
     35:   PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-n', n, flg, ierr))
     36:   nz = 3*n - 4

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

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

     44:   do i = 1, n
     45:     b(i) = 1.0
     46:   end do

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

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

     74:   PetscCallA(MatCreateSeqAIJWithArrays(PETSC_COMM_SELF, n, n, ia, ja, a, J, ierr))
     75:   PetscCallA(MatCreateSeqAIJFromTriple(PETSC_COMM_SELF, n, n, rows, cols, a, Jt, nz, PETSC_FALSE, ierr))
     76:   PetscCallA(MatEqual(J, Jt, equal, ierr))
     77:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jt must be equal')
     78:   PetscCallA(MatDestroy(Jt, ierr))
     79:   PetscCallA(MatCreate(PETSC_COMM_SELF, Jr, ierr))
     80:   PetscCallA(MatSetSizes(Jr, n, n, n, n, ierr))
     81:   PetscCallA(MatSetType(Jr, MATSEQAIJ, ierr))
     82:   PetscCallA(MatSetPreallocationCOO(Jr, nz, rows, cols, ierr))
     83:   PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
     84:   PetscCallA(MatEqual(J, Jr, equal, ierr))
     85:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr must be equal')

     87:   PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, one, n, b, rhs, ierr))
     88:   PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, one, n, x, solution, ierr))

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

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

    100: !     Keep the same size and nonzero structure of the matrix but change its numerical entries
    101:   do i = 2, n - 1
    102:     a(2 + 3*(i - 2) + 1) = 4.0
    103:   end do
    104:   PetscCallA(PetscObjectStateIncrease(J, ierr))
    105:   PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
    106:   PetscCallA(MatEqual(J, Jr, equal, ierr))
    107:   PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr should be equal')
    108:   PetscCallA(MatDestroy(Jr, ierr))

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

    112:   PetscCallA(KSPDestroy(ksp, ierr))
    113:   PetscCallA(VecDestroy(rhs, ierr))
    114:   PetscCallA(VecDestroy(solution, ierr))
    115:   PetscCallA(MatDestroy(J, ierr))

    117:   DEALLOCATE (b, x)
    118:   DEALLOCATE (ia, ja, a)
    119:   DEALLOCATE (rows, cols)

    121:   PetscCallA(PetscFinalize(ierr))
    122: end

    124: !/*TEST
    125: !
    126: !     test:
    127: !       args: -ksp_monitor -ksp_view
    128: !
    129: !TEST*/