1 static const char help[] = "Test MatGetLocalSubMatrix() with multiple levels of nesting.\n\n"; 2 3 #include <petscmat.h> 4 5 int main(int argc, char *argv[]) 6 { 7 IS is0a, is0b, is0, is1, isl0a, isl0b, isl0, isl1; 8 Mat A, Aexplicit; 9 PetscBool usenest, test_mat_is; 10 PetscMPIInt rank, size; 11 PetscInt i, j; 12 13 PetscFunctionBeginUser; 14 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 15 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 16 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 17 18 usenest = PETSC_FALSE; 19 test_mat_is = PETSC_FALSE; 20 PetscCall(PetscOptionsGetBool(NULL, NULL, "-nest", &usenest, NULL)); 21 PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_mat_is", &test_mat_is, NULL)); 22 23 { 24 PetscInt ix0a[1], ix0b[1], ix0[2], ix1[1]; 25 26 ix0a[0] = rank * 2 + 0; 27 ix0b[0] = rank * 2 + 1; 28 ix0[0] = rank * 3 + 0; 29 ix0[1] = rank * 3 + 1; 30 ix1[0] = rank * 3 + 2; 31 PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 1, ix0a, PETSC_COPY_VALUES, &is0a)); 32 PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 1, ix0b, PETSC_COPY_VALUES, &is0b)); 33 PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 2, ix0, PETSC_COPY_VALUES, &is0)); 34 PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 1, ix1, PETSC_COPY_VALUES, &is1)); 35 } 36 { 37 PetscCall(ISCreateStride(PETSC_COMM_SELF, 6, 0, 1, &isl0)); 38 PetscCall(ISCreateStride(PETSC_COMM_SELF, 3, 0, 1, &isl0a)); 39 PetscCall(ISCreateStride(PETSC_COMM_SELF, 3, 3, 1, &isl0b)); 40 PetscCall(ISCreateStride(PETSC_COMM_SELF, 3, 6, 1, &isl1)); 41 } 42 43 if (usenest) { 44 ISLocalToGlobalMapping l2g; 45 PetscInt l2gind[3]; 46 Mat B[9]; 47 48 l2gind[0] = (rank - 1 + size) % size; 49 l2gind[1] = rank; 50 l2gind[2] = (rank + 1) % size; 51 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 3, l2gind, PETSC_COPY_VALUES, &l2g)); 52 for (i = 0; i < 9; i++) { 53 if (test_mat_is) { 54 PetscCall(MatCreateIS(PETSC_COMM_WORLD, 1, 1, 1, PETSC_DECIDE, PETSC_DECIDE, l2g, l2g, &B[i])); 55 } else { 56 PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, 1, 1, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, NULL, PETSC_DECIDE, NULL, &B[i])); 57 PetscCall(MatSetUp(B[i])); 58 PetscCall(MatSetLocalToGlobalMapping(B[i], l2g, l2g)); 59 } 60 } 61 { 62 IS isx[2]; 63 Mat Bx00[4], Bx01[2], Bx10[2]; 64 Mat B00, B01, B10; 65 66 isx[0] = is0a; 67 isx[1] = is0b; 68 Bx00[0] = B[0]; 69 Bx00[1] = B[1]; 70 Bx00[2] = B[3]; 71 Bx00[3] = B[4]; 72 Bx01[0] = B[2]; 73 Bx01[1] = B[5]; 74 Bx10[0] = B[6]; 75 Bx10[1] = B[7]; 76 77 PetscCall(MatCreateNest(PETSC_COMM_WORLD, 2, isx, 2, isx, Bx00, &B00)); 78 PetscCall(MatSetUp(B00)); 79 PetscCall(MatCreateNest(PETSC_COMM_WORLD, 2, isx, 1, NULL, Bx01, &B01)); 80 PetscCall(MatSetUp(B01)); 81 PetscCall(MatCreateNest(PETSC_COMM_WORLD, 1, NULL, 2, isx, Bx10, &B10)); 82 PetscCall(MatSetUp(B10)); 83 { 84 Mat By[4]; 85 IS isy[2]; 86 87 By[0] = B00; 88 By[1] = B01; 89 By[2] = B10; 90 By[3] = B[8]; 91 isy[0] = is0; 92 isy[1] = is1; 93 94 PetscCall(MatCreateNest(PETSC_COMM_WORLD, 2, isy, 2, isy, By, &A)); 95 PetscCall(MatSetUp(A)); 96 } 97 PetscCall(MatDestroy(&B00)); 98 PetscCall(MatDestroy(&B01)); 99 PetscCall(MatDestroy(&B10)); 100 } 101 for (i = 0; i < 9; i++) PetscCall(MatDestroy(&B[i])); 102 PetscCall(ISLocalToGlobalMappingDestroy(&l2g)); 103 } else { 104 ISLocalToGlobalMapping l2g; 105 PetscInt l2gind[9]; 106 for (i = 0; i < 3; i++) 107 for (j = 0; j < 3; j++) l2gind[3 * i + j] = ((rank - 1 + j + size) % size) * 3 + i; 108 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 9, l2gind, PETSC_COPY_VALUES, &l2g)); 109 if (test_mat_is) { 110 PetscCall(MatCreateIS(PETSC_COMM_WORLD, 1, 3, 3, PETSC_DECIDE, PETSC_DECIDE, l2g, l2g, &A)); 111 } else { 112 PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, 3, 3, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, NULL, PETSC_DECIDE, NULL, &A)); 113 PetscCall(MatSetLocalToGlobalMapping(A, l2g, l2g)); 114 } 115 PetscCall(ISLocalToGlobalMappingDestroy(&l2g)); 116 } 117 118 { 119 Mat A00, A11, A0a0a, A0a0b; 120 PetscCall(MatGetLocalSubMatrix(A, isl0, isl0, &A00)); 121 PetscCall(MatGetLocalSubMatrix(A, isl1, isl1, &A11)); 122 PetscCall(MatGetLocalSubMatrix(A00, isl0a, isl0a, &A0a0a)); 123 PetscCall(MatGetLocalSubMatrix(A00, isl0a, isl0b, &A0a0b)); 124 125 PetscCall(MatSetValueLocal(A0a0a, 0, 0, 100 * rank + 1, ADD_VALUES)); 126 PetscCall(MatSetValueLocal(A0a0a, 0, 1, 100 * rank + 2, ADD_VALUES)); 127 PetscCall(MatSetValueLocal(A0a0a, 2, 2, 100 * rank + 9, ADD_VALUES)); 128 129 PetscCall(MatSetValueLocal(A0a0b, 1, 1, 100 * rank + 50 + 5, ADD_VALUES)); 130 131 PetscCall(MatSetValueLocal(A11, 0, 0, 1000 * (rank + 1) + 1, ADD_VALUES)); 132 PetscCall(MatSetValueLocal(A11, 1, 2, 1000 * (rank + 1) + 6, ADD_VALUES)); 133 134 PetscCall(MatRestoreLocalSubMatrix(A00, isl0a, isl0a, &A0a0a)); 135 PetscCall(MatRestoreLocalSubMatrix(A00, isl0a, isl0b, &A0a0b)); 136 PetscCall(MatRestoreLocalSubMatrix(A, isl0, isl0, &A00)); 137 PetscCall(MatRestoreLocalSubMatrix(A, isl1, isl1, &A11)); 138 } 139 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 140 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 141 PetscCall(MatComputeOperator(A, MATAIJ, &Aexplicit)); 142 PetscCall(MatView(Aexplicit, PETSC_VIEWER_STDOUT_WORLD)); 143 PetscCall(MatDestroy(&Aexplicit)); 144 145 { 146 Mat A00, A0a0a, A0a0b; 147 PetscInt rows[] = {0, 1}; 148 PetscCall(MatGetLocalSubMatrix(A, isl0, isl0, &A00)); 149 PetscCall(MatGetLocalSubMatrix(A00, isl0a, isl0a, &A0a0a)); 150 PetscCall(MatGetLocalSubMatrix(A00, isl0a, isl0b, &A0a0b)); 151 152 PetscCall(MatZeroRowsColumnsLocal(A0a0a, 2, rows, 4.0, NULL, NULL)); 153 PetscCall(MatZeroRowsLocal(A0a0b, 1, rows + 1, 0.0, NULL, NULL)); 154 155 PetscCall(MatRestoreLocalSubMatrix(A00, isl0a, isl0a, &A0a0a)); 156 PetscCall(MatRestoreLocalSubMatrix(A00, isl0a, isl0b, &A0a0b)); 157 PetscCall(MatRestoreLocalSubMatrix(A, isl0, isl0, &A00)); 158 } 159 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 160 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 161 PetscCall(MatComputeOperator(A, MATAIJ, &Aexplicit)); 162 PetscCall(MatView(Aexplicit, PETSC_VIEWER_STDOUT_WORLD)); 163 PetscCall(MatDestroy(&Aexplicit)); 164 165 PetscCall(MatDestroy(&A)); 166 PetscCall(ISDestroy(&is0a)); 167 PetscCall(ISDestroy(&is0b)); 168 PetscCall(ISDestroy(&is0)); 169 PetscCall(ISDestroy(&is1)); 170 PetscCall(ISDestroy(&isl0a)); 171 PetscCall(ISDestroy(&isl0b)); 172 PetscCall(ISDestroy(&isl0)); 173 PetscCall(ISDestroy(&isl1)); 174 PetscCall(PetscFinalize()); 175 return 0; 176 } 177 178 /*TEST 179 180 test: 181 nsize: 3 182 args: -test_mat_is {{0 1}} 183 diff_args: -j 184 185 test: 186 suffix: nest 187 nsize: 3 188 args: -nest -test_mat_is {{0 1}} 189 diff_args: -j 190 191 TEST*/ 192