Actual source code: ex261.c
1: static const char help[] = "Tests MatGetDiagonal().\n\n";
3: #include <petscmat.h>
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;
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: }
23: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
24: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
25: PetscCall(MatViewFromOptions(A, NULL, "-mat_view_assembled"));
27: PetscCall(VecSetRandom(diag, NULL));
28: PetscCall(MatGetDiagonal(A, diag));
29: PetscCall(VecViewFromOptions(diag, NULL, "-diag_vec_view"));
31: PetscCall(VecGetLocalSize(diag, &n));
32: PetscCall(VecGetArrayRead(diag, &arr));
33: for (PetscInt i = 0; i < n; ++i) {
34: const PetscScalar lhs = arr[i];
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);
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: }
50: static PetscErrorCode InitializeMatrix(Mat A)
51: {
52: MatType mtype;
53: char *tmp;
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;
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: }
78: PetscCall(MatSetRandom(A, NULL));
79: PetscCall(MatViewFromOptions(A, NULL, "-mat_view_setup"));
80: PetscFunctionReturn(PETSC_SUCCESS);
81: }
83: int main(int argc, char **argv)
84: {
85: Mat A;
86: Vec diag;
88: PetscFunctionBeginUser;
89: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
91: PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
92: PetscCall(MatSetSizes(A, 10, 10, PETSC_DECIDE, PETSC_DECIDE));
93: PetscCall(MatSetFromOptions(A));
95: PetscCall(InitializeMatrix(A));
97: PetscCall(MatCreateVecs(A, &diag, NULL));
99: PetscCall(CheckDiagonal(A, diag, 0.0));
100: PetscCall(CheckDiagonal(A, diag, 1.0));
101: PetscCall(CheckDiagonal(A, diag, 2.0));
103: PetscCall(VecDestroy(&diag));
104: PetscCall(MatDestroy(&A));
105: PetscCall(PetscFinalize());
106: return 0;
107: }
109: /*TEST
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
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
147: TEST*/