1 static const char help[] = "Tests MatGetDiagonal().\n\n"; 2 3 #include <petscmat.h> 4 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 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 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