Actual source code: ex159.c

  1: static const char help[] = "Test MatGetLocalSubMatrix() with multiple levels of nesting.\n\n";

  3: #include <petscmat.h>

  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;

 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));

 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));

 23:   {
 24:     PetscInt ix0a[1], ix0b[1], ix0[2], ix1[1];

 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:   }

 43:   if (usenest) {
 44:     ISLocalToGlobalMapping l2g;
 45:     PetscInt               l2gind[3];
 46:     Mat                    B[9];

 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;

 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];

 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];

 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;

 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:   }

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));

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));

129:     PetscCall(MatSetValueLocal(A0a0b, 1, 1, 100 * rank + 50 + 5, ADD_VALUES));

131:     PetscCall(MatSetValueLocal(A11, 0, 0, 1000 * (rank + 1) + 1, ADD_VALUES));
132:     PetscCall(MatSetValueLocal(A11, 1, 2, 1000 * (rank + 1) + 6, ADD_VALUES));

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));

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));

152:     PetscCall(MatZeroRowsColumnsLocal(A0a0a, 2, rows, 4.0, NULL, NULL));
153:     PetscCall(MatZeroRowsLocal(A0a0b, 1, rows + 1, 0.0, NULL, NULL));

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));

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: }

178: /*TEST

180:    test:
181:       nsize: 3
182:       args: -test_mat_is {{0 1}}
183:       diff_args: -j

185:    test:
186:       suffix: nest
187:       nsize: 3
188:       args: -nest -test_mat_is {{0 1}}
189:       diff_args: -j

191: TEST*/