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,one
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))
32: one = 1
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: PetscCheckA(equal .eqv. PETSC_TRUE,PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jt must be equal')
77: PetscCallA(MatDestroy(Jt,ierr))
78: PetscCallA(MatCreate(PETSC_COMM_SELF,Jr,ierr))
79: PetscCallA(MatSetSizes(Jr,n,n,n,n,ierr))
80: PetscCallA(MatSetType(Jr,MATSEQAIJ,ierr))
81: PetscCallA(MatSetPreallocationCOO(Jr,nz,rows,cols,ierr))
82: PetscCallA(MatSetValuesCOO(Jr,a,INSERT_VALUES,ierr))
83: PetscCallA(MatEqual(J,Jr,equal,ierr))
84: PetscCheckA(equal .eqv. PETSC_TRUE,PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jr must be equal')
86: PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF,one,n,b,rhs,ierr))
87: PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF,one,n,x,solution,ierr))
89: PetscCallA(KSPCreate(PETSC_COMM_SELF,ksp,ierr))
90: PetscCallA(KSPSetErrorIfNotConverged(ksp,PETSC_TRUE,ierr))
91: ! Default to a direct sparse LU solver for robustness
92: PetscCallA(KSPGetPC(ksp,pc,ierr))
93: PetscCallA(PCSetType(pc,PCLU,ierr))
94: PetscCallA(KSPSetFromOptions(ksp,ierr))
95: PetscCallA(KSPSetOperators(ksp,J,J,ierr))
97: PetscCallA(KSPSolve(ksp,rhs,solution,ierr))
99: ! Keep the same size and nonzero structure of the matrix but change its numerical entries
100: do i=2,n-1
101: a(2+3*(i-2)+1) = 4.0;
102: enddo
103: PetscCallA(PetscObjectStateIncrease(J,ierr))
104: PetscCallA(MatSetValuesCOO(Jr,a,INSERT_VALUES,ierr))
105: PetscCallA(MatEqual(J,Jr,equal,ierr))
106: PetscCheckA(equal .eqv. PETSC_TRUE,PETSC_COMM_SELF,PETSC_ERR_PLIB,'Matrices J and Jr should be equal')
107: PetscCallA(MatDestroy(Jr,ierr))
109: PetscCallA(KSPSolve(ksp,rhs,solution,ierr))
111: PetscCallA(KSPDestroy(ksp,ierr))
112: PetscCallA(VecDestroy(rhs,ierr))
113: PetscCallA(VecDestroy(solution,ierr))
114: PetscCallA(MatDestroy(J,ierr))
116: DEALLOCATE (b,x)
117: DEALLOCATE (ia,ja,a)
118: DEALLOCATE (rows,cols)
120: PetscCallA(PetscFinalize(ierr))
121: end
123: !/*TEST
124: !
125: ! test:
126: ! args: -ksp_monitor -ksp_view
127: !
128: !TEST*/