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