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