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*/