xref: /petsc/src/dm/impls/stag/stagmulti.c (revision 1ed6e3ff8437baa411029a28c2b08f047df9ad9a) !
1 /* Internal and DMStag-specific functions related to multigrid */
2 #include <petsc/private/dmstagimpl.h> /*I  "petscdmstag.h"   I*/
3 
4 /*@
5   DMStagRestrictSimple - restricts data from a fine to a coarse `DMSTAG`, in the simplest way
6 
7   Values on coarse cells are averages of all fine cells that they cover.
8   Thus, values on vertices are injected, values on edges are averages
9   of the underlying two fine edges, and values on elements in
10   d dimensions are averages of $2^d$ underlying elements.
11 
12   Input Parameters:
13 + dmf - fine `DM`
14 . xf  - data on fine `DM`
15 - dmc - coarse `DM`
16 
17   Output Parameter:
18 . xc - data on coarse `DM`
19 
20   Level: advanced
21 
22 .seealso: [](ch_stag), `DMSTAG`, `DM`, `DMRestrict()`, `DMCoarsen()`, `DMCreateInjection()`
23 @*/
DMStagRestrictSimple(DM dmf,Vec xf,DM dmc,Vec xc)24 PetscErrorCode DMStagRestrictSimple(DM dmf, Vec xf, DM dmc, Vec xc)
25 {
26   PetscInt dim;
27 
28   PetscFunctionBegin;
29   PetscCall(DMGetDimension(dmf, &dim));
30   if (PetscDefined(USE_DEBUG)) {
31     PetscInt d, nf[DMSTAG_MAX_DIM], nc[DMSTAG_MAX_DIM], doff[DMSTAG_MAX_STRATA], dofc[DMSTAG_MAX_STRATA];
32 
33     PetscCall(DMStagGetLocalSizes(dmf, &nf[0], &nf[1], &nf[2]));
34     PetscCall(DMStagGetLocalSizes(dmc, &nc[0], &nc[1], &nc[2]));
35     for (d = 0; d < dim; ++d) PetscCheck(nf[d] % nc[d] == 0, PetscObjectComm((PetscObject)dmf), PETSC_ERR_PLIB, "Not implemented for non-integer refinement factor");
36     PetscCall(DMStagGetDOF(dmf, &doff[0], &doff[1], &doff[2], &doff[3]));
37     PetscCall(DMStagGetDOF(dmc, &dofc[0], &dofc[1], &dofc[2], &dofc[3]));
38     for (d = 0; d < dim + 1; ++d) PetscCheck(doff[d] == dofc[d], PetscObjectComm((PetscObject)dmf), PETSC_ERR_PLIB, "Cannot transfer between DMStag objects with different dof on each stratum");
39     {
40       PetscInt size_local, entries_local;
41 
42       PetscCall(DMStagGetEntriesLocal(dmf, &entries_local));
43       PetscCall(VecGetLocalSize(xf, &size_local));
44       PetscCheck(entries_local == size_local, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "Fine vector must be a local vector of size %" PetscInt_FMT ", but a vector of size %" PetscInt_FMT " was supplied", entries_local, size_local);
45     }
46     {
47       PetscInt size_local, entries_local;
48 
49       PetscCall(DMStagGetEntriesLocal(dmc, &entries_local));
50       PetscCall(VecGetLocalSize(xc, &size_local));
51       PetscCheck(entries_local == size_local, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "Coarse vector must be a local vector of size %" PetscInt_FMT ", but a vector of size %" PetscInt_FMT " was supplied", entries_local, size_local);
52     }
53   }
54   switch (dim) {
55   case 1:
56     PetscCall(DMStagRestrictSimple_1d(dmf, xf, dmc, xc));
57     break;
58   case 2:
59     PetscCall(DMStagRestrictSimple_2d(dmf, xf, dmc, xc));
60     break;
61   case 3:
62     PetscCall(DMStagRestrictSimple_3d(dmf, xf, dmc, xc));
63     break;
64   default:
65     SETERRQ(PetscObjectComm((PetscObject)dmf), PETSC_ERR_ARG_OUTOFRANGE, "Unsupported dimension %" PetscInt_FMT, dim);
66     break;
67   }
68   PetscFunctionReturn(PETSC_SUCCESS);
69 }
70 
SetInterpolationCoefficientVertex_Private(PetscInt index,PetscInt factor,PetscScalar * a)71 static PetscErrorCode SetInterpolationCoefficientVertex_Private(PetscInt index, PetscInt factor, PetscScalar *a)
72 {
73   PetscFunctionBegin;
74   *a = (index % factor) / (PetscScalar)factor;
75   PetscFunctionReturn(PETSC_SUCCESS);
76 }
77 
SetInterpolationCoefficientCenter_Private(PetscBool belowHalf,PetscInt index,PetscInt factor,PetscScalar * a)78 static PetscErrorCode SetInterpolationCoefficientCenter_Private(PetscBool belowHalf, PetscInt index, PetscInt factor, PetscScalar *a)
79 {
80   PetscFunctionBegin;
81   if (belowHalf) *a = 0.5 + ((index % factor) + 0.5) / (PetscScalar)factor;
82   else *a = ((index % factor) + 0.5) / (PetscScalar)factor - 0.5;
83   PetscFunctionReturn(PETSC_SUCCESS);
84 }
85 
SetInterpolationWeight1d_Private(PetscScalar ax,PetscScalar weight[])86 static PetscErrorCode SetInterpolationWeight1d_Private(PetscScalar ax, PetscScalar weight[])
87 {
88   PetscFunctionBegin;
89   weight[0] = 1.0 - ax;
90   weight[1] = ax;
91   PetscFunctionReturn(PETSC_SUCCESS);
92 }
93 
SetInterpolationWeight2d_Private(PetscScalar ax,PetscScalar ay,PetscScalar weight[])94 static PetscErrorCode SetInterpolationWeight2d_Private(PetscScalar ax, PetscScalar ay, PetscScalar weight[])
95 {
96   PetscFunctionBegin;
97   weight[0] = (1.0 - ax) * (1.0 - ay);
98   weight[1] = ax * (1.0 - ay);
99   weight[2] = (1.0 - ax) * ay;
100   weight[3] = ax * ay;
101   PetscFunctionReturn(PETSC_SUCCESS);
102 }
103 
MergeInterpolationWeightToLeft2d_Private(PetscScalar weight[])104 static PetscErrorCode MergeInterpolationWeightToLeft2d_Private(PetscScalar weight[])
105 {
106   PetscFunctionBegin;
107   weight[0] += weight[1];
108   weight[2] += weight[3];
109   weight[1] = weight[3] = 0.0;
110   PetscFunctionReturn(PETSC_SUCCESS);
111 }
112 
MergeInterpolationWeightToRight2d_Private(PetscScalar weight[])113 static PetscErrorCode MergeInterpolationWeightToRight2d_Private(PetscScalar weight[])
114 {
115   PetscFunctionBegin;
116   weight[1] += weight[0];
117   weight[3] += weight[2];
118   weight[0] = weight[2] = 0.0;
119   PetscFunctionReturn(PETSC_SUCCESS);
120 }
121 
MergeInterpolationWeightToBottom2d_Private(PetscScalar weight[])122 static PetscErrorCode MergeInterpolationWeightToBottom2d_Private(PetscScalar weight[])
123 {
124   PetscFunctionBegin;
125   weight[0] += weight[2];
126   weight[1] += weight[3];
127   weight[2] = weight[3] = 0.0;
128   PetscFunctionReturn(PETSC_SUCCESS);
129 }
130 
MergeInterpolationWeightToTop2d_Private(PetscScalar weight[])131 static PetscErrorCode MergeInterpolationWeightToTop2d_Private(PetscScalar weight[])
132 {
133   PetscFunctionBegin;
134   weight[2] += weight[0];
135   weight[3] += weight[1];
136   weight[0] = weight[1] = 0.0;
137   PetscFunctionReturn(PETSC_SUCCESS);
138 }
139 
SetInterpolationWeight3d_Private(PetscScalar ax,PetscScalar ay,PetscScalar az,PetscScalar weight[])140 static PetscErrorCode SetInterpolationWeight3d_Private(PetscScalar ax, PetscScalar ay, PetscScalar az, PetscScalar weight[])
141 {
142   PetscFunctionBegin;
143   weight[0] = (1.0 - ax) * (1.0 - ay) * (1.0 - az);
144   weight[1] = ax * (1.0 - ay) * (1.0 - az);
145   weight[2] = (1.0 - ax) * ay * (1.0 - az);
146   weight[3] = ax * ay * (1.0 - az);
147   weight[4] = (1.0 - ax) * (1.0 - ay) * az;
148   weight[5] = ax * (1.0 - ay) * az;
149   weight[6] = (1.0 - ax) * ay * az;
150   weight[7] = ax * ay * az;
151   PetscFunctionReturn(PETSC_SUCCESS);
152 }
153 
MergeInterpolationWeightToLeft3d_Private(PetscScalar weight[])154 static PetscErrorCode MergeInterpolationWeightToLeft3d_Private(PetscScalar weight[])
155 {
156   PetscFunctionBegin;
157   weight[0] += weight[1];
158   weight[2] += weight[3];
159   weight[4] += weight[5];
160   weight[6] += weight[7];
161   weight[1] = weight[3] = weight[5] = weight[7] = 0.0;
162   PetscFunctionReturn(PETSC_SUCCESS);
163 }
164 
MergeInterpolationWeightToRight3d_Private(PetscScalar weight[])165 static PetscErrorCode MergeInterpolationWeightToRight3d_Private(PetscScalar weight[])
166 {
167   PetscFunctionBegin;
168   weight[1] += weight[0];
169   weight[3] += weight[2];
170   weight[5] += weight[4];
171   weight[7] += weight[6];
172   weight[0] = weight[2] = weight[4] = weight[6] = 0.0;
173   PetscFunctionReturn(PETSC_SUCCESS);
174 }
175 
MergeInterplationWeightToBottom3d_Private(PetscScalar weight[])176 static PetscErrorCode MergeInterplationWeightToBottom3d_Private(PetscScalar weight[])
177 {
178   PetscFunctionBegin;
179   weight[0] += weight[2];
180   weight[1] += weight[3];
181   weight[4] += weight[6];
182   weight[5] += weight[7];
183   weight[2] = weight[3] = weight[6] = weight[7] = 0.0;
184   PetscFunctionReturn(PETSC_SUCCESS);
185 }
186 
MergeInterpolationWeightToTop3d_Private(PetscScalar weight[])187 static PetscErrorCode MergeInterpolationWeightToTop3d_Private(PetscScalar weight[])
188 {
189   PetscFunctionBegin;
190   weight[2] += weight[0];
191   weight[3] += weight[1];
192   weight[6] += weight[4];
193   weight[7] += weight[5];
194   weight[0] = weight[1] = weight[4] = weight[5] = 0.0;
195   PetscFunctionReturn(PETSC_SUCCESS);
196 }
197 
MergeInterpolationWeightToBack3d_Private(PetscScalar weight[])198 static PetscErrorCode MergeInterpolationWeightToBack3d_Private(PetscScalar weight[])
199 {
200   PetscFunctionBegin;
201   weight[0] += weight[4];
202   weight[1] += weight[5];
203   weight[2] += weight[6];
204   weight[3] += weight[7];
205   weight[4] = weight[5] = weight[6] = weight[7] = 0.0;
206   PetscFunctionReturn(PETSC_SUCCESS);
207 }
208 
MergeInterpolationWeightToFront3d_Private(PetscScalar weight[])209 static PetscErrorCode MergeInterpolationWeightToFront3d_Private(PetscScalar weight[])
210 {
211   PetscFunctionBegin;
212   weight[4] += weight[0];
213   weight[5] += weight[1];
214   weight[6] += weight[2];
215   weight[7] += weight[3];
216   weight[0] = weight[1] = weight[2] = weight[3] = 0.0;
217   PetscFunctionReturn(PETSC_SUCCESS);
218 }
219 
SetRestrictionCoefficientVertex_Private(PetscInt index,PetscInt factor,PetscScalar * a)220 static PetscErrorCode SetRestrictionCoefficientVertex_Private(PetscInt index, PetscInt factor, PetscScalar *a)
221 {
222   PetscFunctionBegin;
223   *a = (factor - PetscAbsInt(index)) / (PetscScalar)(factor * factor);
224   PetscFunctionReturn(PETSC_SUCCESS);
225 }
226 
SetRestrictionCoefficientCenter_Private(PetscInt index,PetscInt factor,PetscScalar * a)227 static PetscErrorCode SetRestrictionCoefficientCenter_Private(PetscInt index, PetscInt factor, PetscScalar *a)
228 {
229   PetscFunctionBegin;
230   if (2 * index + 1 < factor) *a = 0.5 + (index + 0.5) / (PetscScalar)factor;
231   else *a = 1.5 - (index + 0.5) / (PetscScalar)factor;
232   /* Normalization depends on whether the restriction factor is even or odd */
233   if (factor % 2 == 0) *a /= 0.75 * factor;
234   else *a /= (3 * factor * factor + 1) / (PetscScalar)(4 * factor);
235   PetscFunctionReturn(PETSC_SUCCESS);
236 }
237 
MergeRestrictionWeightToLeft1d_Private(PetscScalar weight[],PetscInt m)238 static PetscErrorCode MergeRestrictionWeightToLeft1d_Private(PetscScalar weight[], PetscInt m)
239 {
240   PetscInt       i;
241   const PetscInt dst = m / 2;
242 
243   PetscFunctionBegin;
244   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
245   for (i = m / 2 + 1; i < m; ++i) {
246     weight[dst] += weight[i];
247     weight[i] = 0.0;
248   }
249   PetscFunctionReturn(PETSC_SUCCESS);
250 }
251 
MergeRestrictionWeightToRight1d_Private(PetscScalar weight[],PetscInt m)252 static PetscErrorCode MergeRestrictionWeightToRight1d_Private(PetscScalar weight[], PetscInt m)
253 {
254   PetscInt       i;
255   const PetscInt dst = m / 2;
256 
257   PetscFunctionBegin;
258   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
259   for (i = 0; i < m / 2; ++i) {
260     weight[dst] += weight[i];
261     weight[i] = 0.0;
262   }
263   PetscFunctionReturn(PETSC_SUCCESS);
264 }
265 
MergeRestrictionWeightToLeft2d_Private(PetscScalar weight[],PetscInt m,PetscInt n)266 static PetscErrorCode MergeRestrictionWeightToLeft2d_Private(PetscScalar weight[], PetscInt m, PetscInt n)
267 {
268   PetscInt i, j, src, dst;
269 
270   PetscFunctionBegin;
271   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
272   for (j = 0; j < n; ++j)
273     for (i = m / 2 + 1; i < m; ++i) {
274       src = m * j + i;
275       dst = m * j + m / 2;
276       weight[dst] += weight[src];
277       weight[src] = 0.0;
278     }
279   PetscFunctionReturn(PETSC_SUCCESS);
280 }
281 
MergeRestrictionWeightToRight2d_Private(PetscScalar weight[],PetscInt m,PetscInt n)282 static PetscErrorCode MergeRestrictionWeightToRight2d_Private(PetscScalar weight[], PetscInt m, PetscInt n)
283 {
284   PetscInt i, j, src, dst;
285 
286   PetscFunctionBegin;
287   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
288   for (j = 0; j < n; ++j)
289     for (i = 0; i < m / 2; ++i) {
290       src = m * j + i;
291       dst = m * j + m / 2;
292       weight[dst] += weight[src];
293       weight[src] = 0.0;
294     }
295   PetscFunctionReturn(PETSC_SUCCESS);
296 }
297 
MergeRestrictionWeightToBottom2d_Private(PetscScalar weight[],PetscInt m,PetscInt n)298 static PetscErrorCode MergeRestrictionWeightToBottom2d_Private(PetscScalar weight[], PetscInt m, PetscInt n)
299 {
300   PetscInt i, j, src, dst;
301 
302   PetscFunctionBegin;
303   PetscCheck(n % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
304   for (j = n / 2 + 1; j < n; ++j)
305     for (i = 0; i < m; ++i) {
306       src = m * j + i;
307       dst = m * (n / 2) + i;
308       weight[dst] += weight[src];
309       weight[src] = 0.0;
310     }
311   PetscFunctionReturn(PETSC_SUCCESS);
312 }
313 
MergeRestrictionWeightToTop2d_Private(PetscScalar weight[],PetscInt m,PetscInt n)314 static PetscErrorCode MergeRestrictionWeightToTop2d_Private(PetscScalar weight[], PetscInt m, PetscInt n)
315 {
316   PetscInt i, j, src, dst;
317 
318   PetscFunctionBegin;
319   PetscCheck(n % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
320   for (j = 0; j < n / 2; ++j)
321     for (i = 0; i < m; ++i) {
322       src = m * j + i;
323       dst = m * (n / 2) + i;
324       weight[dst] += weight[src];
325       weight[src] = 0.0;
326     }
327   PetscFunctionReturn(PETSC_SUCCESS);
328 }
329 
MergeRestrictionWeightToLeft3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)330 static PetscErrorCode MergeRestrictionWeightToLeft3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
331 {
332   PetscInt i, j, k, src, dst;
333 
334   PetscFunctionBegin;
335   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
336   for (k = 0; k < p; ++k)
337     for (j = 0; j < n; ++j)
338       for (i = m / 2 + 1; i < m; ++i) {
339         src = m * n * k + m * j + i;
340         dst = m * n * k + m * j + m / 2;
341         weight[dst] += weight[src];
342         weight[src] = 0.0;
343       }
344   PetscFunctionReturn(PETSC_SUCCESS);
345 }
346 
MergeRestrictionWeightToRight3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)347 static PetscErrorCode MergeRestrictionWeightToRight3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
348 {
349   PetscInt i, j, k, src, dst;
350 
351   PetscFunctionBegin;
352   PetscCheck(m % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
353   for (k = 0; k < p; ++k)
354     for (j = 0; j < n; ++j)
355       for (i = 0; i < m / 2; ++i) {
356         src = m * n * k + m * j + i;
357         dst = m * n * k + m * j + m / 2;
358         weight[dst] += weight[src];
359         weight[src] = 0.0;
360       }
361   PetscFunctionReturn(PETSC_SUCCESS);
362 }
363 
MergeRestrictionWeightToBottom3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)364 static PetscErrorCode MergeRestrictionWeightToBottom3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
365 {
366   PetscInt i, j, k, src, dst;
367 
368   PetscFunctionBegin;
369   PetscCheck(n % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
370   for (k = 0; k < p; ++k)
371     for (j = n / 2 + 1; j < n; ++j)
372       for (i = 0; i < m; ++i) {
373         src = m * n * k + m * j + i;
374         dst = m * n * k + m * (n / 2) + i;
375         weight[dst] += weight[src];
376         weight[src] = 0.0;
377       }
378   PetscFunctionReturn(PETSC_SUCCESS);
379 }
380 
MergeRestrictionWeightToTop3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)381 static PetscErrorCode MergeRestrictionWeightToTop3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
382 {
383   PetscInt i, j, k, src, dst;
384 
385   PetscFunctionBegin;
386   PetscCheck(n % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
387   for (k = 0; k < p; ++k)
388     for (j = 0; j < n / 2; ++j)
389       for (i = 0; i < m; ++i) {
390         src = m * n * k + m * j + i;
391         dst = m * n * k + m * (n / 2) + i;
392         weight[dst] += weight[src];
393         weight[src] = 0.0;
394       }
395   PetscFunctionReturn(PETSC_SUCCESS);
396 }
397 
MergeRestrictionWeightToBack3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)398 static PetscErrorCode MergeRestrictionWeightToBack3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
399 {
400   PetscInt i, j, k, src, dst;
401 
402   PetscFunctionBegin;
403   PetscCheck(p % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
404   for (k = p / 2 + 1; k < p; ++k)
405     for (j = 0; j < n; ++j)
406       for (i = 0; i < m; ++i) {
407         src = m * n * k + m * j + i;
408         dst = m * n * (p / 2) + m * j + i;
409         weight[dst] += weight[src];
410         weight[src] = 0.0;
411       }
412   PetscFunctionReturn(PETSC_SUCCESS);
413 }
414 
MergeRestrictionWeightToFront3d_Private(PetscScalar weight[],PetscInt m,PetscInt n,PetscInt p)415 static PetscErrorCode MergeRestrictionWeightToFront3d_Private(PetscScalar weight[], PetscInt m, PetscInt n, PetscInt p)
416 {
417   PetscInt i, j, k, src, dst;
418 
419   PetscFunctionBegin;
420   PetscCheck(p % 2 == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dimension must be odd");
421   for (k = 0; k < p / 2; ++k)
422     for (j = 0; j < n; ++j)
423       for (i = 0; i < m; ++i) {
424         src = m * n * k + m * j + i;
425         dst = m * n * (p / 2) + m * j + i;
426         weight[dst] += weight[src];
427         weight[src] = 0.0;
428       }
429   PetscFunctionReturn(PETSC_SUCCESS);
430 }
431 
RemoveZeroWeights_Private(PetscInt n,DMStagStencil colc[],PetscScalar weight[],PetscInt * count)432 static PetscErrorCode RemoveZeroWeights_Private(PetscInt n, DMStagStencil colc[], PetscScalar weight[], PetscInt *count)
433 {
434   PetscInt i;
435 
436   PetscFunctionBegin;
437   *count = 0;
438   for (i = 0; i < n; ++i)
439     if (weight[i] != 0.0) {
440       colc[*count]   = colc[i];
441       weight[*count] = weight[i];
442       ++(*count);
443     }
444   PetscFunctionReturn(PETSC_SUCCESS);
445 }
446 
447 /* Code duplication note: the next two functions are nearly identical, save the inclusion of the element terms */
DMStagPopulateInterpolation1d_Internal(DM dmc,DM dmf,Mat A)448 PETSC_INTERN PetscErrorCode DMStagPopulateInterpolation1d_Internal(DM dmc, DM dmf, Mat A)
449 {
450   PetscInt       Mc, Mf, factorx, dof[2];
451   PetscInt       xf, mf, nExtraxf, i, d, count;
452   DMStagStencil  rowf, colc[2];
453   PetscScalar    ax, weight[2];
454   PetscInt       ir, ic[2];
455   const PetscInt dim = 1;
456 
457   PetscFunctionBegin;
458   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, NULL, NULL));
459   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, NULL, NULL));
460   factorx = Mf / Mc;
461   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], NULL, NULL));
462 
463   /* In 1D, each fine point can receive data from at most 2 coarse points, at most one of which could be off-process */
464   PetscCall(MatSeqAIJSetPreallocation(A, 2, NULL));
465   PetscCall(MatMPIAIJSetPreallocation(A, 2, NULL, 1, NULL));
466 
467   PetscCall(DMStagGetCorners(dmf, &xf, NULL, NULL, &mf, NULL, NULL, &nExtraxf, NULL, NULL));
468 
469   /* Linear interpolation for vertices */
470   for (d = 0; d < dof[0]; ++d)
471     for (i = xf; i < xf + mf + nExtraxf; ++i) {
472       rowf.i   = i;
473       rowf.c   = d;
474       rowf.loc = DMSTAG_LEFT;
475       for (count = 0; count < 2; ++count) {
476         colc[count].i = i / factorx;
477         colc[count].c = d;
478       }
479       colc[0].loc = DMSTAG_LEFT;
480       colc[1].loc = DMSTAG_RIGHT;
481       PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
482       PetscCall(SetInterpolationWeight1d_Private(ax, weight));
483 
484       PetscCall(RemoveZeroWeights_Private(2, colc, weight, &count));
485       PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
486       PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
487       PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
488     }
489 
490   /* Nearest neighbor for elements */
491   for (d = 0; d < dof[1]; ++d)
492     for (i = xf; i < xf + mf; ++i) {
493       rowf.i      = i;
494       rowf.c      = d;
495       rowf.loc    = DMSTAG_ELEMENT;
496       colc[0].i   = i / factorx;
497       colc[0].c   = d;
498       colc[0].loc = DMSTAG_ELEMENT;
499       weight[0]   = 1.0;
500       PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
501       PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, colc, ic));
502       PetscCall(MatSetValuesLocal(A, 1, &ir, 1, ic, weight, INSERT_VALUES));
503     }
504   PetscFunctionReturn(PETSC_SUCCESS);
505 }
506 
DMStagPopulateInterpolation2d_Internal(DM dmc,DM dmf,Mat A)507 PETSC_INTERN PetscErrorCode DMStagPopulateInterpolation2d_Internal(DM dmc, DM dmf, Mat A)
508 {
509   PetscInt       Mc, Nc, Mf, Nf, factorx, factory, dof[3];
510   PetscInt       xf, yf, mf, nf, nExtraxf, nExtrayf, i, j, d, count;
511   DMStagStencil  rowf, colc[4];
512   PetscScalar    ax, ay, weight[4];
513   PetscInt       ir, ic[4];
514   const PetscInt dim = 2;
515 
516   PetscFunctionBegin;
517   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, &Nc, NULL));
518   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, &Nf, NULL));
519   factorx = Mf / Mc;
520   factory = Nf / Nc;
521   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], &dof[2], NULL));
522 
523   /* In 2D, each fine point can receive data from at most 4 coarse points, at most 3 of which could be off-process */
524   PetscCall(MatSeqAIJSetPreallocation(A, 4, NULL));
525   PetscCall(MatMPIAIJSetPreallocation(A, 4, NULL, 3, NULL));
526 
527   PetscCall(DMStagGetCorners(dmf, &xf, &yf, NULL, &mf, &nf, NULL, &nExtraxf, &nExtrayf, NULL));
528 
529   /* Linear interpolation for vertices */
530   for (d = 0; d < dof[0]; ++d)
531     for (j = yf; j < yf + nf + nExtrayf; ++j)
532       for (i = xf; i < xf + mf + nExtraxf; ++i) {
533         rowf.i   = i;
534         rowf.j   = j;
535         rowf.c   = d;
536         rowf.loc = DMSTAG_DOWN_LEFT;
537         for (count = 0; count < 4; ++count) {
538           colc[count].i = i / factorx;
539           colc[count].j = j / factory;
540           colc[count].c = d;
541         }
542         colc[0].loc = DMSTAG_DOWN_LEFT;
543         colc[1].loc = DMSTAG_DOWN_RIGHT;
544         colc[2].loc = DMSTAG_UP_LEFT;
545         colc[3].loc = DMSTAG_UP_RIGHT;
546         PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
547         PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
548         PetscCall(SetInterpolationWeight2d_Private(ax, ay, weight));
549 
550         PetscCall(RemoveZeroWeights_Private(4, colc, weight, &count));
551         PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
552         PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
553         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
554       }
555 
556   /* Linear interpolation for left edges */
557   for (d = 0; d < dof[1]; ++d)
558     for (j = yf; j < yf + nf; ++j)
559       for (i = xf; i < xf + mf + nExtraxf; ++i) {
560         const PetscBool belowHalfy = (PetscBool)(2 * (j % factory) + 1 < factory);
561 
562         rowf.i   = i;
563         rowf.j   = j;
564         rowf.c   = d;
565         rowf.loc = DMSTAG_LEFT;
566         for (count = 0; count < 4; ++count) {
567           colc[count].i = i / factorx;
568           colc[count].j = j / factory;
569           if (belowHalfy) colc[count].j -= 1;
570           if (count / 2 == 1) colc[count].j += 1;
571           colc[count].c = d;
572         }
573         colc[0].loc = DMSTAG_LEFT;
574         colc[1].loc = DMSTAG_RIGHT;
575         colc[2].loc = DMSTAG_LEFT;
576         colc[3].loc = DMSTAG_RIGHT;
577         PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
578         PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfy, j, factory, &ay));
579         PetscCall(SetInterpolationWeight2d_Private(ax, ay, weight));
580         /* Assume Neumann boundary condition */
581         if (j / factory == 0 && belowHalfy) PetscCall(MergeInterpolationWeightToTop2d_Private(weight));
582         else if (j / factory == Nc - 1 && !belowHalfy) PetscCall(MergeInterpolationWeightToBottom2d_Private(weight));
583 
584         PetscCall(RemoveZeroWeights_Private(4, colc, weight, &count));
585         PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
586         PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
587         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
588       }
589 
590   /* Linear interpolation for down edges */
591   for (d = 0; d < dof[1]; ++d)
592     for (j = yf; j < yf + nf + nExtrayf; ++j)
593       for (i = xf; i < xf + mf; ++i) {
594         const PetscBool belowHalfx = (PetscBool)(2 * (i % factorx) + 1 < factorx);
595 
596         rowf.i   = i;
597         rowf.j   = j;
598         rowf.c   = d;
599         rowf.loc = DMSTAG_DOWN;
600         for (count = 0; count < 4; ++count) {
601           colc[count].i = i / factorx;
602           colc[count].j = j / factory;
603           if (belowHalfx) colc[count].i -= 1;
604           if (count % 2 == 1) colc[count].i += 1;
605           colc[count].c = d;
606         }
607         colc[0].loc = DMSTAG_DOWN;
608         colc[1].loc = DMSTAG_DOWN;
609         colc[2].loc = DMSTAG_UP;
610         colc[3].loc = DMSTAG_UP;
611         PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfx, i, factorx, &ax));
612         PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
613         PetscCall(SetInterpolationWeight2d_Private(ax, ay, weight));
614         /* Assume Neumann boundary condition */
615         if (i / factorx == 0 && belowHalfx) PetscCall(MergeInterpolationWeightToRight2d_Private(weight));
616         else if (i / factorx == Mc - 1 && !belowHalfx) PetscCall(MergeInterpolationWeightToLeft2d_Private(weight));
617 
618         PetscCall(RemoveZeroWeights_Private(4, colc, weight, &count));
619         PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
620         PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
621         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
622       }
623 
624   /* Nearest neighbor for elements */
625   for (d = 0; d < dof[2]; ++d)
626     for (j = yf; j < yf + nf; ++j)
627       for (i = xf; i < xf + mf; ++i) {
628         rowf.i      = i;
629         rowf.j      = j;
630         rowf.c      = d;
631         rowf.loc    = DMSTAG_ELEMENT;
632         colc[0].i   = i / factorx;
633         colc[0].j   = j / factory;
634         colc[0].c   = d;
635         colc[0].loc = DMSTAG_ELEMENT;
636         weight[0]   = 1.0;
637         PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
638         PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, colc, ic));
639         PetscCall(MatSetValuesLocal(A, 1, &ir, 1, ic, weight, INSERT_VALUES));
640       }
641   PetscFunctionReturn(PETSC_SUCCESS);
642 }
643 
DMStagPopulateInterpolation3d_Internal(DM dmc,DM dmf,Mat A)644 PETSC_INTERN PetscErrorCode DMStagPopulateInterpolation3d_Internal(DM dmc, DM dmf, Mat A)
645 {
646   PetscInt       Mc, Nc, Pc, Mf, Nf, Pf, factorx, factory, factorz, dof[4];
647   PetscInt       xf, yf, zf, mf, nf, pf, nExtraxf, nExtrayf, nExtrazf, i, j, k, d, count;
648   DMStagStencil  rowf, colc[8];
649   PetscScalar    ax, ay, az, weight[8];
650   PetscInt       ir, ic[8];
651   const PetscInt dim = 3;
652 
653   PetscFunctionBegin;
654   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, &Nc, &Pc));
655   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, &Nf, &Pf));
656   factorx = Mf / Mc;
657   factory = Nf / Nc;
658   factorz = Pf / Pc;
659   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], &dof[2], &dof[3]));
660 
661   /* In 3D, each fine point can receive data from at most 8 coarse points, at most 7 of which could be off-process */
662   PetscCall(MatSeqAIJSetPreallocation(A, 8, NULL));
663   PetscCall(MatMPIAIJSetPreallocation(A, 8, NULL, 7, NULL));
664 
665   PetscCall(DMStagGetCorners(dmf, &xf, &yf, &zf, &mf, &nf, &pf, &nExtraxf, &nExtrayf, &nExtrazf));
666 
667   /* Linear interpolation for vertices */
668   for (d = 0; d < dof[0]; ++d)
669     for (k = zf; k < zf + pf + nExtrazf; ++k)
670       for (j = yf; j < yf + nf + nExtrayf; ++j)
671         for (i = xf; i < xf + mf + nExtraxf; ++i) {
672           rowf.i   = i;
673           rowf.j   = j;
674           rowf.k   = k;
675           rowf.c   = d;
676           rowf.loc = DMSTAG_BACK_DOWN_LEFT;
677           for (count = 0; count < 8; ++count) {
678             colc[count].i = i / factorx;
679             colc[count].j = j / factory;
680             colc[count].k = k / factorz;
681             colc[count].c = d;
682           }
683           colc[0].loc = DMSTAG_BACK_DOWN_LEFT;
684           colc[1].loc = DMSTAG_BACK_DOWN_RIGHT;
685           colc[2].loc = DMSTAG_BACK_UP_LEFT;
686           colc[3].loc = DMSTAG_BACK_UP_RIGHT;
687           colc[4].loc = DMSTAG_FRONT_DOWN_LEFT;
688           colc[5].loc = DMSTAG_FRONT_DOWN_RIGHT;
689           colc[6].loc = DMSTAG_FRONT_UP_LEFT;
690           colc[7].loc = DMSTAG_FRONT_UP_RIGHT;
691           PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
692           PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
693           PetscCall(SetInterpolationCoefficientVertex_Private(k, factorz, &az));
694           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
695 
696           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
697           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
698           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
699           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
700         }
701 
702   /* Linear interpolation for down left edges */
703   for (d = 0; d < dof[1]; ++d)
704     for (k = zf; k < zf + pf; ++k)
705       for (j = yf; j < yf + nf + nExtrayf; ++j)
706         for (i = xf; i < xf + mf + nExtraxf; ++i) {
707           const PetscBool belowHalfz = (PetscBool)(2 * (k % factorz) + 1 < factorz);
708 
709           rowf.i   = i;
710           rowf.j   = j;
711           rowf.k   = k;
712           rowf.c   = d;
713           rowf.loc = DMSTAG_DOWN_LEFT;
714           for (count = 0; count < 8; ++count) {
715             colc[count].i = i / factorx;
716             colc[count].j = j / factory;
717             colc[count].k = k / factorz;
718             if (belowHalfz) colc[count].k -= 1;
719             if (count / 4 == 1) colc[count].k += 1;
720             colc[count].c = d;
721           }
722           colc[0].loc = DMSTAG_DOWN_LEFT;
723           colc[1].loc = DMSTAG_DOWN_RIGHT;
724           colc[2].loc = DMSTAG_UP_LEFT;
725           colc[3].loc = DMSTAG_UP_RIGHT;
726           colc[4].loc = DMSTAG_DOWN_LEFT;
727           colc[5].loc = DMSTAG_DOWN_RIGHT;
728           colc[6].loc = DMSTAG_UP_LEFT;
729           colc[7].loc = DMSTAG_UP_RIGHT;
730           PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
731           PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
732           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfz, k, factorz, &az));
733           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
734           /* Assume Neumann boundary condition */
735           if (k / factorz == 0 && belowHalfz) PetscCall(MergeInterpolationWeightToFront3d_Private(weight));
736           else if (k / factorz == Pc - 1 && !belowHalfz) PetscCall(MergeInterpolationWeightToBack3d_Private(weight));
737 
738           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
739           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
740           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
741           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
742         }
743 
744   /* Linear interpolation for back left edges */
745   for (d = 0; d < dof[1]; ++d)
746     for (k = zf; k < zf + pf + nExtrazf; ++k)
747       for (j = yf; j < yf + nf; ++j)
748         for (i = xf; i < xf + mf + nExtraxf; ++i) {
749           const PetscBool belowHalfy = (PetscBool)(2 * (j % factory) + 1 < factory);
750 
751           rowf.i   = i;
752           rowf.j   = j;
753           rowf.k   = k;
754           rowf.c   = d;
755           rowf.loc = DMSTAG_BACK_LEFT;
756           for (count = 0; count < 8; ++count) {
757             colc[count].i = i / factorx;
758             colc[count].j = j / factory;
759             colc[count].k = k / factorz;
760             if (belowHalfy) colc[count].j -= 1;
761             if ((count % 4) / 2 == 1) colc[count].j += 1;
762             colc[count].c = d;
763           }
764           colc[0].loc = DMSTAG_BACK_LEFT;
765           colc[1].loc = DMSTAG_BACK_RIGHT;
766           colc[2].loc = DMSTAG_BACK_LEFT;
767           colc[3].loc = DMSTAG_BACK_RIGHT;
768           colc[4].loc = DMSTAG_FRONT_LEFT;
769           colc[5].loc = DMSTAG_FRONT_RIGHT;
770           colc[6].loc = DMSTAG_FRONT_LEFT;
771           colc[7].loc = DMSTAG_FRONT_RIGHT;
772           PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
773           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfy, j, factory, &ay));
774           PetscCall(SetInterpolationCoefficientVertex_Private(k, factorz, &az));
775           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
776           /* Assume Neumann boundary condition */
777           if (j / factory == 0 && belowHalfy) PetscCall(MergeInterpolationWeightToTop3d_Private(weight));
778           else if (j / factory == Nc - 1 && !belowHalfy) PetscCall(MergeInterplationWeightToBottom3d_Private(weight));
779 
780           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
781           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
782           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
783           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
784         }
785 
786   /* Linear interpolation for down back edges */
787   for (d = 0; d < dof[1]; ++d)
788     for (k = zf; k < zf + pf + nExtrazf; ++k)
789       for (j = yf; j < yf + nf + nExtrayf; ++j)
790         for (i = xf; i < xf + mf; ++i) {
791           const PetscBool belowHalfx = (PetscBool)(2 * (i % factorx) + 1 < factorx);
792 
793           rowf.i   = i;
794           rowf.j   = j;
795           rowf.k   = k;
796           rowf.c   = d;
797           rowf.loc = DMSTAG_BACK_DOWN;
798           for (count = 0; count < 8; ++count) {
799             colc[count].i = i / factorx;
800             colc[count].j = j / factory;
801             colc[count].k = k / factorz;
802             if (belowHalfx) colc[count].i -= 1;
803             if (count % 2 == 1) colc[count].i += 1;
804             colc[count].c = d;
805           }
806           colc[0].loc = DMSTAG_BACK_DOWN;
807           colc[1].loc = DMSTAG_BACK_DOWN;
808           colc[2].loc = DMSTAG_BACK_UP;
809           colc[3].loc = DMSTAG_BACK_UP;
810           colc[4].loc = DMSTAG_FRONT_DOWN;
811           colc[5].loc = DMSTAG_FRONT_DOWN;
812           colc[6].loc = DMSTAG_FRONT_UP;
813           colc[7].loc = DMSTAG_FRONT_UP;
814           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfx, i, factorx, &ax));
815           PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
816           PetscCall(SetInterpolationCoefficientVertex_Private(k, factorz, &az));
817           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
818           /* Assume Neumann boundary condition */
819           if (i / factorx == 0 && belowHalfx) PetscCall(MergeInterpolationWeightToRight3d_Private(weight));
820           else if (i / factorx == Mc - 1 && !belowHalfx) PetscCall(MergeInterpolationWeightToLeft3d_Private(weight));
821 
822           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
823           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
824           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
825           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
826         }
827 
828   /* Linear interpolation for left faces */
829   for (d = 0; d < dof[2]; ++d)
830     for (k = zf; k < zf + pf; ++k)
831       for (j = yf; j < yf + nf; ++j)
832         for (i = xf; i < xf + mf + nExtraxf; ++i) {
833           const PetscBool belowHalfy = (PetscBool)(2 * (j % factory) + 1 < factory);
834           const PetscBool belowHalfz = (PetscBool)(2 * (k % factorz) + 1 < factorz);
835 
836           rowf.i   = i;
837           rowf.j   = j;
838           rowf.k   = k;
839           rowf.c   = d;
840           rowf.loc = DMSTAG_LEFT;
841           for (count = 0; count < 8; ++count) {
842             colc[count].i = i / factorx;
843             colc[count].j = j / factory;
844             colc[count].k = k / factorz;
845             if (belowHalfy) colc[count].j -= 1;
846             if ((count % 4) / 2 == 1) colc[count].j += 1;
847             if (belowHalfz) colc[count].k -= 1;
848             if (count / 4 == 1) colc[count].k += 1;
849             colc[count].c = d;
850           }
851           colc[0].loc = DMSTAG_LEFT;
852           colc[1].loc = DMSTAG_RIGHT;
853           colc[2].loc = DMSTAG_LEFT;
854           colc[3].loc = DMSTAG_RIGHT;
855           colc[4].loc = DMSTAG_LEFT;
856           colc[5].loc = DMSTAG_RIGHT;
857           colc[6].loc = DMSTAG_LEFT;
858           colc[7].loc = DMSTAG_RIGHT;
859           PetscCall(SetInterpolationCoefficientVertex_Private(i, factorx, &ax));
860           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfy, j, factory, &ay));
861           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfz, k, factorz, &az));
862           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
863           /* Assume Neumann boundary condition */
864           if (j / factory == 0 && belowHalfy) PetscCall(MergeInterpolationWeightToTop3d_Private(weight));
865           else if (j / factory == Nc - 1 && !belowHalfy) PetscCall(MergeInterplationWeightToBottom3d_Private(weight));
866           if (k / factorz == 0 && belowHalfz) PetscCall(MergeInterpolationWeightToFront3d_Private(weight));
867           else if (k / factorz == Pc - 1 && !belowHalfz) PetscCall(MergeInterpolationWeightToBack3d_Private(weight));
868 
869           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
870           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
871           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
872           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
873         }
874 
875   /* Linear interpolation for down faces */
876   for (d = 0; d < dof[2]; ++d)
877     for (k = zf; k < zf + pf; ++k)
878       for (j = yf; j < yf + nf + nExtrayf; ++j)
879         for (i = xf; i < xf + mf; ++i) {
880           const PetscBool belowHalfx = (PetscBool)(2 * (i % factorx) + 1 < factorx);
881           const PetscBool belowHalfz = (PetscBool)(2 * (k % factorz) + 1 < factorz);
882 
883           rowf.i   = i;
884           rowf.j   = j;
885           rowf.k   = k;
886           rowf.c   = d;
887           rowf.loc = DMSTAG_DOWN;
888           for (count = 0; count < 8; ++count) {
889             colc[count].i = i / factorx;
890             colc[count].j = j / factory;
891             colc[count].k = k / factorz;
892             if (belowHalfx) colc[count].i -= 1;
893             if (count % 2 == 1) colc[count].i += 1;
894             if (belowHalfz) colc[count].k -= 1;
895             if (count / 4 == 1) colc[count].k += 1;
896             colc[count].c = d;
897           }
898           colc[0].loc = DMSTAG_DOWN;
899           colc[1].loc = DMSTAG_DOWN;
900           colc[2].loc = DMSTAG_UP;
901           colc[3].loc = DMSTAG_UP;
902           colc[4].loc = DMSTAG_DOWN;
903           colc[5].loc = DMSTAG_DOWN;
904           colc[6].loc = DMSTAG_UP;
905           colc[7].loc = DMSTAG_UP;
906           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfx, i, factorx, &ax));
907           PetscCall(SetInterpolationCoefficientVertex_Private(j, factory, &ay));
908           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfz, k, factorz, &az));
909           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
910           /* Assume Neumann boundary condition */
911           if (i / factorx == 0 && belowHalfx) PetscCall(MergeInterpolationWeightToRight3d_Private(weight));
912           else if (i / factorx == Mc - 1 && !belowHalfx) PetscCall(MergeInterpolationWeightToLeft3d_Private(weight));
913           if (k / factorz == 0 && belowHalfz) PetscCall(MergeInterpolationWeightToFront3d_Private(weight));
914           else if (k / factorz == Pc - 1 && !belowHalfz) PetscCall(MergeInterpolationWeightToBack3d_Private(weight));
915 
916           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
917           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
918           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
919           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
920         }
921 
922   /* Linear interpolation for back faces */
923   for (d = 0; d < dof[2]; ++d)
924     for (k = zf; k < zf + pf + nExtrazf; ++k)
925       for (j = yf; j < yf + nf; ++j)
926         for (i = xf; i < xf + mf; ++i) {
927           const PetscBool belowHalfx = (PetscBool)(2 * (i % factorx) + 1 < factorx);
928           const PetscBool belowHalfy = (PetscBool)(2 * (j % factory) + 1 < factory);
929 
930           rowf.i   = i;
931           rowf.j   = j;
932           rowf.k   = k;
933           rowf.c   = d;
934           rowf.loc = DMSTAG_BACK;
935           for (count = 0; count < 8; ++count) {
936             colc[count].i = i / factorx;
937             colc[count].j = j / factory;
938             colc[count].k = k / factorz;
939             if (belowHalfx) colc[count].i -= 1;
940             if (count % 2 == 1) colc[count].i += 1;
941             if (belowHalfy) colc[count].j -= 1;
942             if ((count % 4) / 2 == 1) colc[count].j += 1;
943             colc[count].c = d;
944           }
945           colc[0].loc = DMSTAG_BACK;
946           colc[1].loc = DMSTAG_BACK;
947           colc[2].loc = DMSTAG_BACK;
948           colc[3].loc = DMSTAG_BACK;
949           colc[4].loc = DMSTAG_FRONT;
950           colc[5].loc = DMSTAG_FRONT;
951           colc[6].loc = DMSTAG_FRONT;
952           colc[7].loc = DMSTAG_FRONT;
953           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfx, i, factorx, &ax));
954           PetscCall(SetInterpolationCoefficientCenter_Private(belowHalfy, j, factory, &ay));
955           PetscCall(SetInterpolationCoefficientVertex_Private(k, factorz, &az));
956           PetscCall(SetInterpolationWeight3d_Private(ax, ay, az, weight));
957           /* Assume Neumann boundary condition */
958           if (i / factorx == 0 && belowHalfx) PetscCall(MergeInterpolationWeightToRight3d_Private(weight));
959           else if (i / factorx == Mc - 1 && !belowHalfx) PetscCall(MergeInterpolationWeightToLeft3d_Private(weight));
960           if (j / factory == 0 && belowHalfy) PetscCall(MergeInterpolationWeightToTop3d_Private(weight));
961           else if (j / factory == Nc - 1 && !belowHalfy) PetscCall(MergeInterplationWeightToBottom3d_Private(weight));
962 
963           PetscCall(RemoveZeroWeights_Private(8, colc, weight, &count));
964           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
965           PetscCall(DMStagStencilToIndexLocal(dmc, dim, count, colc, ic));
966           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
967         }
968 
969   /* Nearest neighbor for elements */
970   for (d = 0; d < dof[3]; ++d)
971     for (k = zf; k < zf + pf; ++k)
972       for (j = yf; j < yf + nf; ++j)
973         for (i = xf; i < xf + mf; ++i) {
974           rowf.i      = i;
975           rowf.j      = j;
976           rowf.k      = k;
977           rowf.c      = d;
978           rowf.loc    = DMSTAG_ELEMENT;
979           colc[0].i   = i / factorx;
980           colc[0].j   = j / factory;
981           colc[0].k   = k / factorz;
982           colc[0].c   = d;
983           colc[0].loc = DMSTAG_ELEMENT;
984           weight[0]   = 1.0;
985           PetscCall(DMStagStencilToIndexLocal(dmf, dim, 1, &rowf, &ir));
986           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, colc, ic));
987           PetscCall(MatSetValuesLocal(A, 1, &ir, 1, ic, weight, INSERT_VALUES));
988         }
989   PetscFunctionReturn(PETSC_SUCCESS);
990 }
991 
DMStagPopulateRestriction1d_Internal(DM dmc,DM dmf,Mat A)992 PETSC_INTERN PetscErrorCode DMStagPopulateRestriction1d_Internal(DM dmc, DM dmf, Mat A)
993 {
994   PetscInt       Mc, Mf, factorx, dof[2];
995   PetscInt       xc, mc, nExtraxc, i, d, ii, count;
996   PetscInt       maxFinePoints, maxOffRankFinePoints;
997   DMStagStencil  rowc;
998   DMStagStencil *colf;
999   PetscScalar   *weight;
1000   PetscInt       ir;
1001   PetscInt      *ic;
1002   const PetscInt dim = 1;
1003 
1004   PetscFunctionBegin;
1005   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, NULL, NULL));
1006   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, NULL, NULL));
1007   factorx = Mf / Mc;
1008   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], NULL, NULL));
1009 
1010   /* In 1D, each coarse point can receive from up to (2 * factorx - 1) fine points, (factorx - 1) of which may be off-rank */
1011   maxFinePoints        = 2 * factorx - 1;
1012   maxOffRankFinePoints = maxFinePoints - factorx;
1013   PetscCall(MatSeqAIJSetPreallocation(A, maxFinePoints, NULL));
1014   PetscCall(MatMPIAIJSetPreallocation(A, maxFinePoints, NULL, maxOffRankFinePoints, NULL));
1015   PetscCall(PetscMalloc3(maxFinePoints, &colf, maxFinePoints, &weight, maxFinePoints, &ic));
1016 
1017   PetscCall(DMStagGetCorners(dmc, &xc, NULL, NULL, &mc, NULL, NULL, &nExtraxc, NULL, NULL));
1018 
1019   for (d = 0; d < dof[0]; ++d)
1020     for (i = xc; i < xc + mc + nExtraxc; ++i) {
1021       rowc.i   = i;
1022       rowc.c   = d;
1023       rowc.loc = DMSTAG_LEFT;
1024       count    = 0;
1025       for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1026         colf[count].i   = i * factorx + ii;
1027         colf[count].c   = d;
1028         colf[count].loc = DMSTAG_LEFT;
1029         PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &weight[count]));
1030         ++count;
1031       }
1032       if (i == 0) PetscCall(MergeRestrictionWeightToRight1d_Private(weight, 2 * factorx - 1));
1033       else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft1d_Private(weight, 2 * factorx - 1));
1034 
1035       PetscCall(RemoveZeroWeights_Private(2 * factorx - 1, colf, weight, &count));
1036       PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1037       PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1038       PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1039     }
1040 
1041   for (d = 0; d < dof[1]; ++d)
1042     for (i = xc; i < xc + mc; ++i) {
1043       rowc.i   = i;
1044       rowc.c   = d;
1045       rowc.loc = DMSTAG_ELEMENT;
1046       count    = 0;
1047       for (ii = 0; ii < factorx; ++ii) {
1048         colf[count].i   = i * factorx + ii;
1049         colf[count].c   = d;
1050         colf[count].loc = DMSTAG_ELEMENT;
1051         weight[count]   = 1 / (PetscScalar)factorx;
1052         ++count;
1053       }
1054 
1055       PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1056       PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1057       PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1058     }
1059 
1060   PetscCall(PetscFree3(colf, weight, ic));
1061   PetscFunctionReturn(PETSC_SUCCESS);
1062 }
1063 
DMStagPopulateRestriction2d_Internal(DM dmc,DM dmf,Mat A)1064 PETSC_INTERN PetscErrorCode DMStagPopulateRestriction2d_Internal(DM dmc, DM dmf, Mat A)
1065 {
1066   PetscInt       Mc, Nc, Mf, Nf, factorx, factory, dof[3];
1067   PetscInt       xc, yc, mc, nc, nExtraxc, nExtrayc, i, j, d, ii, jj, count;
1068   PetscInt       maxFinePoints, maxOffRankFinePoints;
1069   DMStagStencil  rowc;
1070   DMStagStencil *colf;
1071   PetscScalar    ax, ay;
1072   PetscScalar   *weight;
1073   PetscInt       ir;
1074   PetscInt      *ic;
1075   const PetscInt dim = 2;
1076 
1077   PetscFunctionBegin;
1078   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, &Nc, NULL));
1079   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, &Nf, NULL));
1080   factorx = Mf / Mc;
1081   factory = Nf / Nc;
1082   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], &dof[2], NULL));
1083 
1084   /* In 2D, each coarse point can receive from up to ((2 * factorx - 1) * (2 * factory - 1)) fine points,
1085      up to ((2 * factorx - 1) * (2 * factory - 1) - factorx * factory) of which may be off rank */
1086   maxFinePoints        = (2 * factorx - 1) * (2 * factory - 1);
1087   maxOffRankFinePoints = maxFinePoints - factorx * factory;
1088   PetscCall(MatSeqAIJSetPreallocation(A, maxFinePoints, NULL));
1089   PetscCall(MatMPIAIJSetPreallocation(A, maxFinePoints, NULL, maxOffRankFinePoints, NULL));
1090   PetscCall(PetscMalloc3(maxFinePoints, &colf, maxFinePoints, &weight, maxFinePoints, &ic));
1091 
1092   PetscCall(DMStagGetCorners(dmc, &xc, &yc, NULL, &mc, &nc, NULL, &nExtraxc, &nExtrayc, NULL));
1093 
1094   for (d = 0; d < dof[0]; ++d)
1095     for (j = yc; j < yc + nc + nExtrayc; ++j)
1096       for (i = xc; i < xc + mc + nExtraxc; ++i) {
1097         rowc.i   = i;
1098         rowc.j   = j;
1099         rowc.c   = d;
1100         rowc.loc = DMSTAG_DOWN_LEFT;
1101         count    = 0;
1102         for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1103           for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1104             colf[count].i   = i * factorx + ii;
1105             colf[count].j   = j * factory + jj;
1106             colf[count].c   = d;
1107             colf[count].loc = DMSTAG_DOWN_LEFT;
1108             PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1109             PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1110             weight[count] = ax * ay;
1111             ++count;
1112           }
1113         if (i == 0) PetscCall(MergeRestrictionWeightToRight2d_Private(weight, 2 * factorx - 1, 2 * factory - 1));
1114         else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft2d_Private(weight, 2 * factorx - 1, 2 * factory - 1));
1115         if (j == 0) PetscCall(MergeRestrictionWeightToTop2d_Private(weight, 2 * factorx - 1, 2 * factory - 1));
1116         else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom2d_Private(weight, 2 * factorx - 1, 2 * factory - 1));
1117 
1118         PetscCall(RemoveZeroWeights_Private((2 * factorx - 1) * (2 * factory - 1), colf, weight, &count));
1119         PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1120         PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1121         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1122       }
1123 
1124   for (d = 0; d < dof[1]; ++d)
1125     for (j = yc; j < yc + nc; ++j)
1126       for (i = xc; i < xc + mc + nExtraxc; ++i) {
1127         rowc.i   = i;
1128         rowc.j   = j;
1129         rowc.c   = d;
1130         rowc.loc = DMSTAG_LEFT;
1131         count    = 0;
1132         for (jj = 0; jj < factory; ++jj)
1133           for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1134             colf[count].i   = i * factorx + ii;
1135             colf[count].j   = j * factory + jj;
1136             colf[count].c   = d;
1137             colf[count].loc = DMSTAG_LEFT;
1138             PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1139             PetscCall(SetRestrictionCoefficientCenter_Private(jj, factory, &ay));
1140             weight[count] = ax * ay;
1141             ++count;
1142           }
1143         if (i == 0) PetscCall(MergeRestrictionWeightToRight2d_Private(weight, 2 * factorx - 1, factory));
1144         else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft2d_Private(weight, 2 * factorx - 1, factory));
1145 
1146         PetscCall(RemoveZeroWeights_Private((2 * factorx - 1) * factory, colf, weight, &count));
1147         PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1148         PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1149         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1150       }
1151 
1152   for (d = 0; d < dof[1]; ++d)
1153     for (j = yc; j < yc + nc + nExtrayc; ++j)
1154       for (i = xc; i < xc + mc; ++i) {
1155         rowc.i   = i;
1156         rowc.j   = j;
1157         rowc.c   = d;
1158         rowc.loc = DMSTAG_DOWN;
1159         count    = 0;
1160         for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1161           for (ii = 0; ii < factorx; ++ii) {
1162             colf[count].i   = i * factorx + ii;
1163             colf[count].j   = j * factory + jj;
1164             colf[count].c   = d;
1165             colf[count].loc = DMSTAG_DOWN;
1166             PetscCall(SetRestrictionCoefficientCenter_Private(ii, factorx, &ax));
1167             PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1168             weight[count] = ax * ay;
1169             ++count;
1170           }
1171         if (j == 0) PetscCall(MergeRestrictionWeightToTop2d_Private(weight, factorx, 2 * factory - 1));
1172         else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom2d_Private(weight, factorx, 2 * factory - 1));
1173 
1174         PetscCall(RemoveZeroWeights_Private(factorx * (2 * factory - 1), colf, weight, &count));
1175         PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1176         PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1177         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1178       }
1179 
1180   for (d = 0; d < dof[2]; ++d)
1181     for (j = yc; j < yc + nc; ++j)
1182       for (i = xc; i < xc + mc; ++i) {
1183         rowc.i   = i;
1184         rowc.j   = j;
1185         rowc.c   = d;
1186         rowc.loc = DMSTAG_ELEMENT;
1187         count    = 0;
1188         for (jj = 0; jj < factory; ++jj)
1189           for (ii = 0; ii < factorx; ++ii) {
1190             colf[count].i   = i * factorx + ii;
1191             colf[count].j   = j * factory + jj;
1192             colf[count].c   = d;
1193             colf[count].loc = DMSTAG_ELEMENT;
1194             weight[count]   = 1 / (PetscScalar)(factorx * factory);
1195             ++count;
1196           }
1197 
1198         PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1199         PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1200         PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1201       }
1202 
1203   PetscCall(PetscFree3(colf, weight, ic));
1204   PetscFunctionReturn(PETSC_SUCCESS);
1205 }
1206 
DMStagPopulateRestriction3d_Internal(DM dmc,DM dmf,Mat A)1207 PETSC_INTERN PetscErrorCode DMStagPopulateRestriction3d_Internal(DM dmc, DM dmf, Mat A)
1208 {
1209   PetscInt       Mc, Nc, Pc, Mf, Nf, Pf, factorx, factory, factorz, dof[4];
1210   PetscInt       xc, yc, zc, mc, nc, pc, nExtraxc, nExtrayc, nExtrazc, i, j, k, d, ii, jj, kk, count;
1211   PetscInt       maxFinePoints, maxOffRankFinePoints;
1212   DMStagStencil  rowc;
1213   DMStagStencil *colf;
1214   PetscScalar    ax, ay, az;
1215   PetscScalar   *weight;
1216   PetscInt       ir;
1217   PetscInt      *ic;
1218   const PetscInt dim = 3;
1219 
1220   PetscFunctionBegin;
1221   PetscCall(DMStagGetGlobalSizes(dmc, &Mc, &Nc, &Pc));
1222   PetscCall(DMStagGetGlobalSizes(dmf, &Mf, &Nf, &Pf));
1223   factorx = Mf / Mc;
1224   factory = Nf / Nc;
1225   factorz = Pf / Pc;
1226   PetscCall(DMStagGetDOF(dmc, &dof[0], &dof[1], &dof[2], &dof[3]));
1227 
1228   /* In 2D, each coarse point can receive from up to ((2 * factorx - 1) * (2 * factory - 1) * (2 * factorz - 1)) fine points,
1229      up to ((2 * factorx - 1) * (2 * factory - 1) * (2 * factorz - 1) - factorx * factory * factorz) of which may be off rank */
1230   maxFinePoints        = (2 * factorx - 1) * (2 * factory - 1) * (2 * factorz - 1);
1231   maxOffRankFinePoints = maxFinePoints - factorx * factory * factorz;
1232   PetscCall(MatSeqAIJSetPreallocation(A, maxFinePoints, NULL));
1233   PetscCall(MatMPIAIJSetPreallocation(A, maxFinePoints, NULL, maxOffRankFinePoints, NULL));
1234   PetscCall(PetscMalloc3(maxFinePoints, &colf, maxFinePoints, &weight, maxFinePoints, &ic));
1235 
1236   PetscCall(DMStagGetCorners(dmc, &xc, &yc, &zc, &mc, &nc, &pc, &nExtraxc, &nExtrayc, &nExtrazc));
1237 
1238   for (d = 0; d < dof[0]; ++d)
1239     for (k = zc; k < zc + pc + nExtrazc; ++k)
1240       for (j = yc; j < yc + nc + nExtrayc; ++j)
1241         for (i = xc; i < xc + mc + nExtraxc; ++i) {
1242           rowc.i   = i;
1243           rowc.j   = j;
1244           rowc.k   = k;
1245           rowc.c   = d;
1246           rowc.loc = DMSTAG_BACK_DOWN_LEFT;
1247           count    = 0;
1248           for (kk = -(factorz - 1); kk <= factorz - 1; ++kk)
1249             for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1250               for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1251                 colf[count].i   = i * factorx + ii;
1252                 colf[count].j   = j * factory + jj;
1253                 colf[count].k   = k * factorz + kk;
1254                 colf[count].c   = d;
1255                 colf[count].loc = DMSTAG_BACK_DOWN_LEFT;
1256                 PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1257                 PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1258                 PetscCall(SetRestrictionCoefficientVertex_Private(kk, factorz, &az));
1259                 weight[count] = ax * ay * az;
1260                 ++count;
1261               }
1262           if (i == 0) PetscCall(MergeRestrictionWeightToRight3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1263           else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1264           if (j == 0) PetscCall(MergeRestrictionWeightToTop3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1265           else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1266           if (k == 0) PetscCall(MergeRestrictionWeightToFront3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1267           else if (k == Pc) PetscCall(MergeRestrictionWeightToBack3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, 2 * factorz - 1));
1268 
1269           PetscCall(RemoveZeroWeights_Private((2 * factorx - 1) * (2 * factory - 1) * (2 * factorz - 1), colf, weight, &count));
1270           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1271           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1272           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1273         }
1274 
1275   for (d = 0; d < dof[1]; ++d)
1276     for (k = zc; k < zc + pc; ++k)
1277       for (j = yc; j < yc + nc + nExtrayc; ++j)
1278         for (i = xc; i < xc + mc + nExtraxc; ++i) {
1279           rowc.i   = i;
1280           rowc.j   = j;
1281           rowc.k   = k;
1282           rowc.c   = d;
1283           rowc.loc = DMSTAG_DOWN_LEFT;
1284           count    = 0;
1285           for (kk = 0; kk < factorz; ++kk)
1286             for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1287               for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1288                 colf[count].i   = i * factorx + ii;
1289                 colf[count].j   = j * factory + jj;
1290                 colf[count].k   = k * factorz + kk;
1291                 colf[count].c   = d;
1292                 colf[count].loc = DMSTAG_DOWN_LEFT;
1293                 PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1294                 PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1295                 PetscCall(SetRestrictionCoefficientCenter_Private(kk, factorz, &az));
1296                 weight[count] = ax * ay * az;
1297                 ++count;
1298               }
1299           if (i == 0) PetscCall(MergeRestrictionWeightToRight3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, factorz));
1300           else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, factorz));
1301           if (j == 0) PetscCall(MergeRestrictionWeightToTop3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, factorz));
1302           else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom3d_Private(weight, 2 * factorx - 1, 2 * factory - 1, factorz));
1303 
1304           PetscCall(RemoveZeroWeights_Private((2 * factorx - 1) * (2 * factory - 1) * factorz, colf, weight, &count));
1305           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1306           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1307           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1308         }
1309 
1310   for (d = 0; d < dof[1]; ++d)
1311     for (k = zc; k < zc + pc + nExtrazc; ++k)
1312       for (j = yc; j < yc + nc; ++j)
1313         for (i = xc; i < xc + mc + nExtraxc; ++i) {
1314           rowc.i   = i;
1315           rowc.j   = j;
1316           rowc.k   = k;
1317           rowc.c   = d;
1318           rowc.loc = DMSTAG_BACK_LEFT;
1319           count    = 0;
1320           for (kk = -(factorz - 1); kk <= factorz - 1; ++kk)
1321             for (jj = 0; jj < factory; ++jj)
1322               for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1323                 colf[count].i   = i * factorx + ii;
1324                 colf[count].j   = j * factory + jj;
1325                 colf[count].k   = k * factorz + kk;
1326                 colf[count].c   = d;
1327                 colf[count].loc = DMSTAG_BACK_LEFT;
1328                 PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1329                 PetscCall(SetRestrictionCoefficientCenter_Private(jj, factory, &ay));
1330                 PetscCall(SetRestrictionCoefficientVertex_Private(kk, factorz, &az));
1331                 weight[count] = ax * ay * az;
1332                 ++count;
1333               }
1334           if (i == 0) PetscCall(MergeRestrictionWeightToRight3d_Private(weight, 2 * factorx - 1, factory, 2 * factorz - 1));
1335           else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft3d_Private(weight, 2 * factorx - 1, factory, 2 * factorz - 1));
1336           if (k == 0) PetscCall(MergeRestrictionWeightToFront3d_Private(weight, 2 * factorx - 1, factory, 2 * factorz - 1));
1337           else if (k == Pc) PetscCall(MergeRestrictionWeightToBack3d_Private(weight, 2 * factorx - 1, factory, 2 * factorz - 1));
1338         }
1339 
1340   for (d = 0; d < dof[1]; ++d)
1341     for (k = zc; k < zc + pc + nExtrazc; ++k)
1342       for (j = yc; j < yc + nc + nExtrayc; ++j)
1343         for (i = xc; i < xc + mc; ++i) {
1344           rowc.i   = i;
1345           rowc.j   = j;
1346           rowc.k   = k;
1347           rowc.c   = d;
1348           rowc.loc = DMSTAG_BACK_DOWN;
1349           count    = 0;
1350           for (kk = -(factorz - 1); kk <= factorz - 1; ++kk)
1351             for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1352               for (ii = 0; ii < factorx; ++ii) {
1353                 colf[count].i   = i * factorx + ii;
1354                 colf[count].j   = j * factory + jj;
1355                 colf[count].k   = k * factorz + kk;
1356                 colf[count].c   = d;
1357                 colf[count].loc = DMSTAG_BACK_DOWN;
1358                 PetscCall(SetRestrictionCoefficientCenter_Private(ii, factorx, &ax));
1359                 PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1360                 PetscCall(SetRestrictionCoefficientVertex_Private(kk, factorz, &az));
1361                 weight[count] = ax * ay * az;
1362                 ++count;
1363               }
1364           if (j == 0) PetscCall(MergeRestrictionWeightToTop3d_Private(weight, factorx, 2 * factory - 1, 2 * factorz - 1));
1365           else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom3d_Private(weight, factorx, 2 * factory - 1, 2 * factorz - 1));
1366           if (k == 0) PetscCall(MergeRestrictionWeightToFront3d_Private(weight, factorx, 2 * factory - 1, 2 * factorz - 1));
1367           else if (k == Pc) PetscCall(MergeRestrictionWeightToBack3d_Private(weight, factorx, 2 * factory - 1, 2 * factorz - 1));
1368 
1369           PetscCall(RemoveZeroWeights_Private(factorx * (2 * factory - 1) * (2 * factorz - 1), colf, weight, &count));
1370           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1371           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1372           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1373         }
1374 
1375   for (d = 0; d < dof[2]; ++d)
1376     for (k = zc; k < zc + pc; ++k)
1377       for (j = yc; j < yc + nc; ++j)
1378         for (i = xc; i < xc + mc + nExtraxc; ++i) {
1379           rowc.i   = i;
1380           rowc.j   = j;
1381           rowc.k   = k;
1382           rowc.c   = d;
1383           rowc.loc = DMSTAG_LEFT;
1384           count    = 0;
1385           for (kk = 0; kk < factorz; ++kk)
1386             for (jj = 0; jj < factory; ++jj)
1387               for (ii = -(factorx - 1); ii <= factorx - 1; ++ii) {
1388                 colf[count].i   = i * factorx + ii;
1389                 colf[count].j   = j * factory + jj;
1390                 colf[count].k   = k * factorz + kk;
1391                 colf[count].c   = d;
1392                 colf[count].loc = DMSTAG_LEFT;
1393                 PetscCall(SetRestrictionCoefficientVertex_Private(ii, factorx, &ax));
1394                 PetscCall(SetRestrictionCoefficientCenter_Private(jj, factory, &ay));
1395                 PetscCall(SetRestrictionCoefficientCenter_Private(kk, factorz, &az));
1396                 weight[count] = ax * ay * az;
1397                 ++count;
1398               }
1399           if (i == 0) PetscCall(MergeRestrictionWeightToRight3d_Private(weight, 2 * factorx - 1, factory, factorz));
1400           else if (i == Mc) PetscCall(MergeRestrictionWeightToLeft3d_Private(weight, 2 * factorx - 1, factory, factorz));
1401 
1402           PetscCall(RemoveZeroWeights_Private((2 * factorx - 1) * factory * factorz, colf, weight, &count));
1403           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1404           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1405           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1406         }
1407 
1408   for (d = 0; d < dof[2]; ++d)
1409     for (k = zc; k < zc + pc; ++k)
1410       for (j = yc; j < yc + nc + nExtrayc; ++j)
1411         for (i = xc; i < xc + mc; ++i) {
1412           rowc.i   = i;
1413           rowc.j   = j;
1414           rowc.k   = k;
1415           rowc.c   = d;
1416           rowc.loc = DMSTAG_DOWN;
1417           count    = 0;
1418           for (kk = 0; kk < factorz; ++kk)
1419             for (jj = -(factory - 1); jj <= factory - 1; ++jj)
1420               for (ii = 0; ii < factorx; ++ii) {
1421                 colf[count].i   = i * factorx + ii;
1422                 colf[count].j   = j * factory + jj;
1423                 colf[count].k   = k * factorz + kk;
1424                 colf[count].c   = d;
1425                 colf[count].loc = DMSTAG_DOWN;
1426                 PetscCall(SetRestrictionCoefficientCenter_Private(ii, factorx, &ax));
1427                 PetscCall(SetRestrictionCoefficientVertex_Private(jj, factory, &ay));
1428                 PetscCall(SetRestrictionCoefficientCenter_Private(kk, factorz, &az));
1429                 weight[count] = ax * ay * az;
1430                 ++count;
1431               }
1432           if (j == 0) PetscCall(MergeRestrictionWeightToTop3d_Private(weight, factorx, 2 * factory - 1, factorz));
1433           else if (j == Nc) PetscCall(MergeRestrictionWeightToBottom3d_Private(weight, factorx, 2 * factory - 1, factorz));
1434 
1435           PetscCall(RemoveZeroWeights_Private(factorx * (2 * factory - 1) * factorz, colf, weight, &count));
1436           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1437           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1438           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1439         }
1440 
1441   for (d = 0; d < dof[2]; ++d)
1442     for (k = zc; k < zc + pc + nExtrazc; ++k)
1443       for (j = yc; j < yc + nc; ++j)
1444         for (i = xc; i < xc + mc; ++i) {
1445           rowc.i   = i;
1446           rowc.j   = j;
1447           rowc.k   = k;
1448           rowc.c   = d;
1449           rowc.loc = DMSTAG_BACK;
1450           count    = 0;
1451           for (kk = -(factorz - 1); kk <= factorz - 1; ++kk)
1452             for (jj = 0; jj < factory; ++jj)
1453               for (ii = 0; ii < factorx; ++ii) {
1454                 colf[count].i   = i * factorx + ii;
1455                 colf[count].j   = j * factory + jj;
1456                 colf[count].k   = k * factorz + kk;
1457                 colf[count].c   = d;
1458                 colf[count].loc = DMSTAG_BACK;
1459                 PetscCall(SetRestrictionCoefficientCenter_Private(ii, factorx, &ax));
1460                 PetscCall(SetRestrictionCoefficientCenter_Private(jj, factory, &ay));
1461                 PetscCall(SetRestrictionCoefficientVertex_Private(kk, factorz, &az));
1462                 weight[count] = ax * ay * az;
1463                 ++count;
1464               }
1465           if (k == 0) PetscCall(MergeRestrictionWeightToFront3d_Private(weight, factorx, factory, 2 * factorz - 1));
1466           else if (k == Pc) PetscCall(MergeRestrictionWeightToBack3d_Private(weight, factorx, factory, 2 * factorz - 1));
1467 
1468           PetscCall(RemoveZeroWeights_Private(factorx * factory * (2 * factorz - 1), colf, weight, &count));
1469           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1470           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1471           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1472         }
1473 
1474   for (d = 0; d < dof[3]; ++d)
1475     for (k = zc; k < zc + pc; ++k)
1476       for (j = yc; j < yc + nc; ++j)
1477         for (i = xc; i < xc + mc; ++i) {
1478           rowc.i   = i;
1479           rowc.j   = j;
1480           rowc.k   = k;
1481           rowc.c   = d;
1482           rowc.loc = DMSTAG_ELEMENT;
1483           count    = 0;
1484           for (kk = 0; kk < factorz; ++kk)
1485             for (jj = 0; jj < factory; ++jj)
1486               for (ii = 0; ii < factorx; ++ii) {
1487                 colf[count].i   = i * factorx + ii;
1488                 colf[count].j   = j * factory + jj;
1489                 colf[count].k   = k * factorz + kk;
1490                 colf[count].c   = d;
1491                 colf[count].loc = DMSTAG_ELEMENT;
1492                 weight[count]   = 1 / (PetscScalar)(factorx * factory * factorz);
1493                 ++count;
1494               }
1495 
1496           PetscCall(DMStagStencilToIndexLocal(dmc, dim, 1, &rowc, &ir));
1497           PetscCall(DMStagStencilToIndexLocal(dmf, dim, count, colf, ic));
1498           PetscCall(MatSetValuesLocal(A, 1, &ir, count, ic, weight, INSERT_VALUES));
1499         }
1500 
1501   PetscCall(PetscFree3(colf, weight, ic));
1502   PetscFunctionReturn(PETSC_SUCCESS);
1503 }
1504