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