1 static char help[] = "Demonstration of flexible submesh assembly.\n\n";
2
3 #include <petscksp.h>
4
main(int argc,char ** argv)5 int main(int argc, char **argv)
6 {
7 Mat G, A, B, C, D;
8 ISLocalToGlobalMapping rowMap, colMap;
9 IS row, col;
10 PetscInt *locRows, *locCols;
11 PetscInt m = 5, n = 5, M = PETSC_DETERMINE, N = PETSC_DETERMINE, rStart, cStart;
12 PetscBool isnest;
13
14 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
15 // Fill up blocks with local lexicographic numbering
16 PetscCall(PetscSplitOwnership(PETSC_COMM_WORLD, &m, &M));
17 PetscCall(PetscSplitOwnership(PETSC_COMM_WORLD, &n, &N));
18 PetscCall(MatCreate(PETSC_COMM_WORLD, &G));
19 PetscCall(MatSetSizes(G, m, n, M, N));
20 PetscCall(MatSetFromOptions(G));
21 PetscCall(PetscObjectTypeCompare((PetscObject)G, MATNEST, &isnest));
22 PetscCall(MatGetOwnershipRange(G, &rStart, NULL));
23 PetscCall(MatGetOwnershipRangeColumn(G, &cStart, NULL));
24 if (isnest) {
25 Mat submat[4];
26 PetscLayout layoutA, layoutD;
27 PetscInt *locRowsA, *locColsA, *locRowsD, *locColsD, rStartA, rStartD;
28 ISLocalToGlobalMapping rowMapA, colMapA, rowMapD, colMapD;
29
30 PetscCall(PetscLayoutCreate(PETSC_COMM_WORLD, &layoutA));
31 PetscCall(PetscLayoutSetLocalSize(layoutA, 3));
32 PetscCall(PetscLayoutSetUp(layoutA));
33 PetscCall(PetscLayoutGetRange(layoutA, &rStartA, NULL));
34 PetscCall(PetscLayoutDestroy(&layoutA));
35 PetscCall(PetscMalloc1(3, &locRowsA));
36 for (PetscInt r = 0; r < 3; ++r) locRowsA[r] = r + rStartA;
37 PetscCall(PetscMalloc1(3, &locColsA));
38 for (PetscInt c = 0; c < 3; ++c) locColsA[c] = c + rStartA;
39 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 3, locRowsA, PETSC_OWN_POINTER, &rowMapA));
40 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 3, locColsA, PETSC_OWN_POINTER, &colMapA));
41 PetscCall(PetscLayoutCreate(PETSC_COMM_WORLD, &layoutD));
42 PetscCall(PetscLayoutSetLocalSize(layoutD, 2));
43 PetscCall(PetscLayoutSetUp(layoutD));
44 PetscCall(PetscLayoutGetRange(layoutD, &rStartD, NULL));
45 PetscCall(PetscLayoutDestroy(&layoutD));
46 PetscCall(PetscMalloc1(2, &locRowsD));
47 for (PetscInt r = 0; r < 2; ++r) locRowsD[r] = r + rStartD;
48 PetscCall(PetscMalloc1(2, &locColsD));
49 for (PetscInt c = 0; c < 2; ++c) locColsD[c] = c + rStartD;
50 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 2, locRowsD, PETSC_OWN_POINTER, &rowMapD));
51 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 2, locColsD, PETSC_OWN_POINTER, &colMapD));
52
53 PetscCall(MatCreate(PETSC_COMM_WORLD, &submat[0]));
54 PetscCall(MatSetSizes(submat[0], 3, 3, PETSC_DETERMINE, PETSC_DETERMINE));
55 PetscCall(MatSetType(submat[0], MATAIJ));
56 PetscCall(MatSetLocalToGlobalMapping(submat[0], rowMapA, colMapA));
57 PetscCall(MatCreate(PETSC_COMM_WORLD, &submat[1]));
58 PetscCall(MatSetSizes(submat[1], 3, 2, PETSC_DETERMINE, PETSC_DETERMINE));
59 PetscCall(MatSetType(submat[1], MATAIJ));
60 PetscCall(MatSetLocalToGlobalMapping(submat[1], rowMapA, colMapD));
61 PetscCall(MatCreate(PETSC_COMM_WORLD, &submat[2]));
62 PetscCall(MatSetSizes(submat[2], 2, 3, PETSC_DETERMINE, PETSC_DETERMINE));
63 PetscCall(MatSetType(submat[2], MATAIJ));
64 PetscCall(MatSetLocalToGlobalMapping(submat[2], rowMapD, colMapA));
65 PetscCall(MatCreate(PETSC_COMM_WORLD, &submat[3]));
66 PetscCall(MatSetSizes(submat[3], 2, 2, PETSC_DETERMINE, PETSC_DETERMINE));
67 PetscCall(MatSetType(submat[3], MATAIJ));
68 PetscCall(MatSetLocalToGlobalMapping(submat[3], rowMapD, colMapD));
69 for (PetscInt i = 0; i < 4; ++i) PetscCall(MatSetUp(submat[i]));
70 PetscCall(MatNestSetSubMats(G, 2, NULL, 2, NULL, submat));
71 for (PetscInt i = 0; i < 4; ++i) PetscCall(MatDestroy(&submat[i]));
72
73 PetscCall(ISLocalToGlobalMappingDestroy(&rowMapA));
74 PetscCall(ISLocalToGlobalMappingDestroy(&colMapA));
75 PetscCall(ISLocalToGlobalMappingDestroy(&rowMapD));
76 PetscCall(ISLocalToGlobalMappingDestroy(&colMapD));
77 }
78 PetscCall(PetscMalloc1(m, &locRows));
79 for (PetscInt r = 0; r < m; ++r) locRows[r] = r + rStart;
80 PetscCall(PetscMalloc1(n, &locCols));
81 for (PetscInt c = 0; c < n; ++c) locCols[c] = c + cStart;
82 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, m, locRows, PETSC_OWN_POINTER, &rowMap));
83 PetscCall(ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, n, locCols, PETSC_OWN_POINTER, &colMap));
84 PetscCall(MatSetLocalToGlobalMapping(G, rowMap, colMap));
85 PetscCall(ISLocalToGlobalMappingDestroy(&rowMap));
86 PetscCall(ISLocalToGlobalMappingDestroy(&colMap));
87
88 // (0,0) Block A
89 PetscInt A_row[] = {0, 1, 2};
90 PetscInt A_col[] = {0, 1, 2};
91
92 m = PETSC_STATIC_ARRAY_LENGTH(A_row);
93 n = PETSC_STATIC_ARRAY_LENGTH(A_col);
94 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, m, A_row, PETSC_COPY_VALUES, &row));
95 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, A_col, PETSC_COPY_VALUES, &col));
96 PetscCall(MatGetLocalSubMatrix(G, row, col, &A));
97 PetscCall(ISDestroy(&row));
98 PetscCall(ISDestroy(&col));
99 for (PetscInt i = 0; i < m; ++i) {
100 for (PetscInt j = 0; j < n; ++j) {
101 PetscScalar v = i * n + j;
102
103 PetscCall(MatSetValuesLocal(A, 1, &i, 1, &j, &v, INSERT_VALUES));
104 }
105 }
106 PetscCall(MatDestroy(&A));
107
108 // (0,1) Block B
109 PetscInt B_row[] = {0, 1, 2};
110 PetscInt B_col[] = {3, 4};
111
112 m = PETSC_STATIC_ARRAY_LENGTH(B_row);
113 n = PETSC_STATIC_ARRAY_LENGTH(B_col);
114 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, m, B_row, PETSC_COPY_VALUES, &row));
115 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, B_col, PETSC_COPY_VALUES, &col));
116 PetscCall(MatGetLocalSubMatrix(G, row, col, &B));
117 PetscCall(ISDestroy(&row));
118 PetscCall(ISDestroy(&col));
119 for (PetscInt i = 0; i < m; ++i) {
120 for (PetscInt j = 0; j < n; ++j) {
121 PetscScalar v = i * n + j;
122
123 PetscCall(MatSetValuesLocal(B, 1, &i, 1, &j, &v, INSERT_VALUES));
124 }
125 }
126 PetscCall(MatDestroy(&B));
127
128 // (0,1) Block C
129 PetscInt C_row[] = {3, 4};
130 PetscInt C_col[] = {0, 1, 2};
131
132 m = PETSC_STATIC_ARRAY_LENGTH(C_row);
133 n = PETSC_STATIC_ARRAY_LENGTH(C_col);
134 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, m, C_row, PETSC_COPY_VALUES, &row));
135 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, C_col, PETSC_COPY_VALUES, &col));
136 PetscCall(MatGetLocalSubMatrix(G, row, col, &C));
137 PetscCall(ISDestroy(&row));
138 PetscCall(ISDestroy(&col));
139 for (PetscInt i = 0; i < m; ++i) {
140 for (PetscInt j = 0; j < n; ++j) {
141 PetscScalar v = i * n + j;
142
143 PetscCall(MatSetValuesLocal(C, 1, &i, 1, &j, &v, INSERT_VALUES));
144 }
145 }
146 PetscCall(MatDestroy(&C));
147
148 // (0,0) Block D
149 PetscInt D_row[] = {3, 4};
150 PetscInt D_col[] = {3, 4};
151
152 m = PETSC_STATIC_ARRAY_LENGTH(D_row);
153 n = PETSC_STATIC_ARRAY_LENGTH(D_col);
154 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, m, D_row, PETSC_COPY_VALUES, &row));
155 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, n, D_col, PETSC_COPY_VALUES, &col));
156 PetscCall(MatGetLocalSubMatrix(G, row, col, &D));
157 PetscCall(ISDestroy(&row));
158 PetscCall(ISDestroy(&col));
159 for (PetscInt i = 0; i < m; ++i) {
160 for (PetscInt j = 0; j < n; ++j) {
161 PetscScalar v = i * n + j;
162
163 PetscCall(MatSetValuesLocal(D, 1, &i, 1, &j, &v, INSERT_VALUES));
164 }
165 }
166 PetscCall(MatDestroy(&D));
167
168 PetscCall(MatAssemblyBegin(G, MAT_FINAL_ASSEMBLY));
169 PetscCall(MatAssemblyEnd(G, MAT_FINAL_ASSEMBLY));
170 PetscCall(MatViewFromOptions(G, NULL, "-G_view"));
171 PetscCall(MatDestroy(&G));
172 PetscCall(PetscFinalize());
173 return 0;
174 }
175
176 /*TEST
177 test:
178 suffix: 0
179 args: -mat_type aij -G_view
180
181 test:
182 suffix: 1
183 args: -mat_type nest -G_view -mat_view_nest_sub
184
185 test:
186 suffix: 2
187 nsize: 2
188 args: -mat_type aij -G_view
189
190 test:
191 suffix: 3
192 nsize: 2
193 args: -mat_type nest -G_view
194 TEST*/
195