xref: /petsc/src/mat/tests/ex159.c (revision 0b46e949f18ac28417071477034640c76a0832a0)
1 static const char help[] = "Test MatGetLocalSubMatrix() with multiple levels of nesting.\n\n";
2 
3 #include <petscmat.h>
4 
main(int argc,char * argv[])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