1 static const char help[] = "Tests MatGetDiagonal().\n\n";
2
3 #include <petscmat.h>
4
CheckDiagonal(Mat A,Vec diag,PetscScalar dval)5 static PetscErrorCode CheckDiagonal(Mat A, Vec diag, PetscScalar dval)
6 {
7 static PetscBool first_time = PETSC_TRUE;
8 const PetscReal rtol = 1e-10, atol = PETSC_SMALL;
9 PetscInt rstart, rend, n;
10 const PetscScalar *arr;
11
12 PetscFunctionBegin;
13 PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
14 // If matrix is AIJ, MatSetRandom() will have randomly chosen the locations of nonzeros,
15 // which may not be on the diagonal. So a reallocation is not necessarily a bad thing here.
16 if (first_time) PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE));
17 for (PetscInt i = rstart; i < rend; ++i) PetscCall(MatSetValue(A, i, i, dval, INSERT_VALUES));
18 if (first_time) {
19 PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE));
20 first_time = PETSC_FALSE;
21 }
22
23 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
24 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
25 PetscCall(MatViewFromOptions(A, NULL, "-mat_view_assembled"));
26
27 PetscCall(VecSetRandom(diag, NULL));
28 PetscCall(MatGetDiagonal(A, diag));
29 PetscCall(VecViewFromOptions(diag, NULL, "-diag_vec_view"));
30
31 PetscCall(VecGetLocalSize(diag, &n));
32 PetscCall(VecGetArrayRead(diag, &arr));
33 for (PetscInt i = 0; i < n; ++i) {
34 const PetscScalar lhs = arr[i];
35
36 if (!PetscIsCloseAtTolScalar(lhs, dval, rtol, atol)) {
37 const PetscReal lhs_r = PetscRealPart(lhs);
38 const PetscReal lhs_i = PetscImaginaryPart(lhs);
39 const PetscReal dval_r = PetscRealPart(dval);
40 const PetscReal dval_i = PetscImaginaryPart(dval);
41
42 PetscCheck(PetscIsCloseAtTol(lhs_r, dval_r, rtol, atol), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Real component actual[%" PetscInt_FMT "] %g != expected[%" PetscInt_FMT "] %g", i, (double)lhs_r, i, (double)dval_r);
43 PetscCheck(PetscIsCloseAtTol(lhs_i, dval_i, rtol, atol), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Imaginary component actual[%" PetscInt_FMT "] %g != expected[%" PetscInt_FMT "] %g", i, (double)lhs_i, i, (double)dval_i);
44 }
45 }
46 PetscCall(VecRestoreArrayRead(diag, &arr));
47 PetscFunctionReturn(PETSC_SUCCESS);
48 }
49
InitializeMatrix(Mat A)50 static PetscErrorCode InitializeMatrix(Mat A)
51 {
52 MatType mtype;
53 char *tmp;
54
55 PetscFunctionBegin;
56 PetscCall(MatSetUp(A));
57 PetscCall(MatGetType(A, &mtype));
58 PetscCall(PetscStrstr(mtype, "aij", &tmp));
59 if (tmp) {
60 PetscInt rows, cols, diag_nnz, offdiag_nnz;
61 PetscInt *dnnz, *onnz;
62
63 // is some kind of AIJ
64 PetscCall(MatGetLocalSize(A, &rows, &cols));
65 // at least 3 nonzeros in diagonal block
66 diag_nnz = PetscMin(cols, 3);
67 // leave at least 3 *zeros* per row
68 offdiag_nnz = PetscMax(cols - diag_nnz - 3, 0);
69 PetscCall(PetscMalloc2(rows, &dnnz, rows, &onnz));
70 for (PetscInt i = 0; i < rows; ++i) {
71 dnnz[i] = diag_nnz;
72 onnz[i] = offdiag_nnz;
73 }
74 PetscCall(MatXAIJSetPreallocation(A, PETSC_DECIDE, dnnz, onnz, NULL, NULL));
75 PetscCall(PetscFree2(dnnz, onnz));
76 }
77
78 PetscCall(MatSetRandom(A, NULL));
79 PetscCall(MatViewFromOptions(A, NULL, "-mat_view_setup"));
80 PetscFunctionReturn(PETSC_SUCCESS);
81 }
82
main(int argc,char ** argv)83 int main(int argc, char **argv)
84 {
85 Mat A;
86 Vec diag;
87
88 PetscFunctionBeginUser;
89 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
90
91 PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
92 PetscCall(MatSetSizes(A, 10, 10, PETSC_DECIDE, PETSC_DECIDE));
93 PetscCall(MatSetFromOptions(A));
94
95 PetscCall(InitializeMatrix(A));
96
97 PetscCall(MatCreateVecs(A, &diag, NULL));
98
99 PetscCall(CheckDiagonal(A, diag, 0.0));
100 PetscCall(CheckDiagonal(A, diag, 1.0));
101 PetscCall(CheckDiagonal(A, diag, 2.0));
102
103 PetscCall(VecDestroy(&diag));
104 PetscCall(MatDestroy(&A));
105 PetscCall(PetscFinalize());
106 return 0;
107 }
108
109 /*TEST
110
111 testset:
112 output_file: output/empty.out
113 nsize: {{1 2}}
114 suffix: dense
115 test:
116 suffix: standard
117 args: -mat_type dense
118 test:
119 suffix: cuda
120 requires: cuda
121 args: -mat_type densecuda
122 test:
123 suffix: hip
124 requires: hip
125 args: -mat_type densehip
126
127 testset:
128 output_file: output/empty.out
129 nsize: {{1 2}}
130 suffix: aij
131 test:
132 suffix: standard
133 args: -mat_type aij
134 test:
135 suffix: cuda
136 requires: cuda
137 args: -mat_type aijcusparse
138 test:
139 suffix: hip
140 requires: hip
141 args: -mat_type aijhipsparse
142 test:
143 suffix: kokkos
144 requires: kokkos kokkos_kernels
145 args: -mat_type aijkokkos
146
147 TEST*/
148