xref: /libCEED/examples/petsc/src/swarmutils.c (revision 78f7fce354c760f980d0580f63d29ea51c63cedc)
1 // Copyright (c) 2017-2022, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 #include "../include/swarmutils.h"
9 #include "../include/matops.h"
10 #include "../qfunctions/swarm/swarmmass.h"
11 
12 // ------------------------------------------------------------------------------------------------
13 // Context utilities
14 // ------------------------------------------------------------------------------------------------
15 PetscErrorCode DMSwarmCeedContextCreate(DM dm_swarm, const char *ceed_resource, DMSwarmCeedContext *ctx) {
16   DM                  dm_mesh, dm_coord;
17   CeedElemRestriction elem_restr_u_mesh, elem_restr_x_mesh, elem_restr_x_points, elem_restr_u_points, elem_restr_q_data_points;
18   CeedBasis           basis_u, basis_x;
19   CeedVector          x_ref_points, q_data_points;
20   CeedInt             num_comp;
21 
22   PetscFunctionBeginUser;
23   PetscCall(PetscNew(ctx));
24   PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
25   PetscCall(DMGetCoordinateDM(dm_mesh, &dm_coord));
26 
27   CeedInit(ceed_resource, &(*ctx)->ceed);
28   // Background mesh objects
29   {
30     BPData bp_data = {.q_mode = CEED_GAUSS};
31 
32     PetscCall(CreateBasisFromPlex((*ctx)->ceed, dm_mesh, NULL, 0, 0, 0, bp_data, &basis_u));
33     PetscCall(CreateBasisFromPlex((*ctx)->ceed, dm_coord, NULL, 0, 0, 0, bp_data, &basis_x));
34     PetscCall(CreateRestrictionFromPlex((*ctx)->ceed, dm_mesh, 0, NULL, 0, &elem_restr_u_mesh));
35     PetscCall(CreateRestrictionFromPlex((*ctx)->ceed, dm_coord, 0, NULL, 0, &elem_restr_x_mesh));
36 
37     // -- Mesh vectors
38     CeedElemRestrictionCreateVector(elem_restr_u_mesh, &(*ctx)->u_mesh, NULL);
39     CeedElemRestrictionCreateVector(elem_restr_u_mesh, &(*ctx)->v_mesh, NULL);
40   }
41   // Swarm objects
42   {
43     PetscInt        dim;
44     const PetscInt *cell_points;
45     IS              is_points;
46     Vec             X_ref;
47     CeedInt         num_elem;
48 
49     PetscCall(DMSwarmCreateReferenceCoordinates(dm_swarm, &is_points, &X_ref));
50     PetscCall(DMGetDimension(dm_mesh, &dim));
51     CeedElemRestrictionGetNumElements(elem_restr_u_mesh, &num_elem);
52     CeedElemRestrictionGetNumComponents(elem_restr_u_mesh, &num_comp);
53 
54     PetscCall(ISGetIndices(is_points, &cell_points));
55     PetscInt num_points = cell_points[num_elem + 1] - num_elem - 2;
56     CeedInt  offsets[num_elem + 1 + num_points];
57 
58     for (PetscInt i = 0; i < num_elem + 1; i++) offsets[i] = cell_points[i + 1] - 1;
59     for (PetscInt i = num_elem + 1; i < num_points + num_elem + 1; i++) offsets[i] = cell_points[i + 1];
60     PetscCall(ISRestoreIndices(is_points, &cell_points));
61 
62     // -- Points restrictions
63     CeedElemRestrictionCreateAtPoints((*ctx)->ceed, num_elem, num_points, num_comp, num_points * num_comp, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
64                                       &elem_restr_u_points);
65     CeedElemRestrictionCreateAtPoints((*ctx)->ceed, num_elem, num_points, dim, num_points * dim, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
66                                       &elem_restr_x_points);
67     CeedElemRestrictionCreateAtPoints((*ctx)->ceed, num_elem, num_points, 1, num_points, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
68                                       &elem_restr_q_data_points);
69 
70     // -- Points vectors
71     CeedElemRestrictionCreateVector(elem_restr_u_points, &(*ctx)->u_points, NULL);
72     CeedElemRestrictionCreateVector(elem_restr_q_data_points, &q_data_points, NULL);
73 
74     // -- Ref coordinates
75     {
76       PetscMemType       X_mem_type;
77       const PetscScalar *x;
78 
79       CeedVectorCreate((*ctx)->ceed, num_points * dim, &x_ref_points);
80 
81       PetscCall(VecGetArrayReadAndMemType(X_ref, (const PetscScalar **)&x, &X_mem_type));
82       CeedVectorSetArray(x_ref_points, MemTypeP2C(X_mem_type), CEED_COPY_VALUES, (CeedScalar *)x);
83       PetscCall(VecRestoreArrayReadAndMemType(X_ref, (const PetscScalar **)&x));
84     }
85 
86     // Create Q data
87     {
88       CeedQFunction qf_setup;
89       CeedOperator  op_setup;
90       CeedVector    x_coord;
91 
92       {
93         Vec                X_loc;
94         CeedInt            len;
95         const PetscScalar *x;
96 
97         PetscCall(DMGetCoordinatesLocal(dm_mesh, &X_loc));
98         PetscCall(VecGetLocalSize(X_loc, &len));
99         CeedVectorCreate((*ctx)->ceed, len, &x_coord);
100 
101         PetscCall(VecGetArrayRead(X_loc, &x));
102         CeedVectorSetArray(x_coord, CEED_MEM_HOST, CEED_COPY_VALUES, (CeedScalar *)x);
103         PetscCall(VecRestoreArrayRead(X_loc, &x));
104       }
105 
106       // Setup geometric scaling
107       CeedQFunctionCreateInterior((*ctx)->ceed, 1, SetupMass, SetupMass_loc, &qf_setup);
108       CeedQFunctionAddInput(qf_setup, "x", dim * dim, CEED_EVAL_GRAD);
109       CeedQFunctionAddInput(qf_setup, "weight", 1, CEED_EVAL_WEIGHT);
110       CeedQFunctionAddOutput(qf_setup, "rho", 1, CEED_EVAL_NONE);
111 
112       CeedOperatorCreateAtPoints((*ctx)->ceed, qf_setup, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE, &op_setup);
113       CeedOperatorSetField(op_setup, "x", elem_restr_x_mesh, basis_x, CEED_VECTOR_ACTIVE);
114       CeedOperatorSetField(op_setup, "weight", CEED_ELEMRESTRICTION_NONE, basis_x, CEED_VECTOR_NONE);
115       CeedOperatorSetField(op_setup, "rho", elem_restr_q_data_points, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE);
116       CeedOperatorAtPointsSetPoints(op_setup, elem_restr_x_points, x_ref_points);
117 
118       CeedOperatorApply(op_setup, x_coord, q_data_points, CEED_REQUEST_IMMEDIATE);
119 
120       // Cleanup
121       CeedVectorDestroy(&x_coord);
122       CeedQFunctionDestroy(&qf_setup);
123       CeedOperatorDestroy(&op_setup);
124     }
125 
126     // -- Cleanup
127     PetscCall(ISDestroy(&is_points));
128     PetscCall(VecDestroy(&X_ref));
129   }
130 
131   PetscCall(DMSetApplicationContext(dm_mesh, (void *)(*ctx)));
132 
133   // Create operators
134   // Mesh to points interpolation operator
135   {
136     CeedQFunction qf_mesh_to_points;
137 
138     // -- Create operator
139     CeedQFunctionCreateIdentity((*ctx)->ceed, num_comp, CEED_EVAL_INTERP, CEED_EVAL_NONE, &qf_mesh_to_points);
140 
141     CeedOperatorCreateAtPoints((*ctx)->ceed, qf_mesh_to_points, NULL, NULL, &(*ctx)->op_mesh_to_points);
142     CeedOperatorSetField((*ctx)->op_mesh_to_points, "input", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
143     CeedOperatorSetField((*ctx)->op_mesh_to_points, "output", elem_restr_u_points, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE);
144     CeedOperatorAtPointsSetPoints((*ctx)->op_mesh_to_points, elem_restr_x_points, x_ref_points);
145 
146     // -- Cleanup
147     CeedQFunctionDestroy(&qf_mesh_to_points);
148   }
149 
150   // RHS operator
151   {
152     CeedQFunction        qf_pts_to_mesh;
153     CeedQFunctionContext qf_ctx;
154 
155     // -- Mass QFunction
156     CeedQFunctionCreateInterior((*ctx)->ceed, 1, Mass, Mass_loc, &qf_pts_to_mesh);
157     CeedQFunctionAddInput(qf_pts_to_mesh, "q data", 1, CEED_EVAL_NONE);
158     CeedQFunctionAddInput(qf_pts_to_mesh, "u", num_comp, CEED_EVAL_NONE);
159     CeedQFunctionAddOutput(qf_pts_to_mesh, "v", num_comp, CEED_EVAL_INTERP);
160 
161     // -- QFunction context
162     CeedQFunctionContextCreate((*ctx)->ceed, &qf_ctx);
163     CeedQFunctionContextSetData(qf_ctx, CEED_MEM_HOST, CEED_COPY_VALUES, sizeof(num_comp), &num_comp);
164     CeedQFunctionSetContext(qf_pts_to_mesh, qf_ctx);
165 
166     // -- Mass Operator
167     CeedOperatorCreateAtPoints((*ctx)->ceed, qf_pts_to_mesh, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE, &(*ctx)->op_points_to_mesh);
168     CeedOperatorSetField((*ctx)->op_points_to_mesh, "q data", elem_restr_q_data_points, CEED_BASIS_NONE, q_data_points);
169     CeedOperatorSetField((*ctx)->op_points_to_mesh, "u", elem_restr_u_points, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE);
170     CeedOperatorSetField((*ctx)->op_points_to_mesh, "v", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
171     CeedOperatorAtPointsSetPoints((*ctx)->op_points_to_mesh, elem_restr_x_points, x_ref_points);
172 
173     // -- Cleanup
174     CeedQFunctionContextDestroy(&qf_ctx);
175     CeedQFunctionDestroy(&qf_pts_to_mesh);
176   }
177 
178   // Mass operator
179   {
180     CeedQFunction        qf_mass;
181     CeedQFunctionContext ctx_mass;
182 
183     // -- Mass QFunction
184     CeedQFunctionCreateInterior((*ctx)->ceed, 1, Mass, Mass_loc, &qf_mass);
185     CeedQFunctionAddInput(qf_mass, "q data", 1, CEED_EVAL_NONE);
186     CeedQFunctionAddInput(qf_mass, "u", num_comp, CEED_EVAL_INTERP);
187     CeedQFunctionAddOutput(qf_mass, "v", num_comp, CEED_EVAL_INTERP);
188 
189     // -- QFunction context
190     CeedQFunctionContextCreate((*ctx)->ceed, &ctx_mass);
191     CeedQFunctionContextSetData(ctx_mass, CEED_MEM_HOST, CEED_COPY_VALUES, sizeof(num_comp), &num_comp);
192     CeedQFunctionSetContext(qf_mass, ctx_mass);
193 
194     // -- Mass Operator
195     CeedOperatorCreateAtPoints((*ctx)->ceed, qf_mass, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE, &(*ctx)->op_mass);
196     CeedOperatorSetField((*ctx)->op_mass, "q data", elem_restr_q_data_points, CEED_BASIS_NONE, q_data_points);
197     CeedOperatorSetField((*ctx)->op_mass, "u", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
198     CeedOperatorSetField((*ctx)->op_mass, "v", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
199     CeedOperatorAtPointsSetPoints((*ctx)->op_mass, elem_restr_x_points, x_ref_points);
200 
201     // -- Cleanup
202     CeedQFunctionContextDestroy(&ctx_mass);
203     CeedQFunctionDestroy(&qf_mass);
204   }
205 
206   // Cleanup
207   CeedElemRestrictionDestroy(&elem_restr_u_mesh);
208   CeedElemRestrictionDestroy(&elem_restr_x_mesh);
209   CeedElemRestrictionDestroy(&elem_restr_u_points);
210   CeedElemRestrictionDestroy(&elem_restr_x_points);
211   CeedElemRestrictionDestroy(&elem_restr_q_data_points);
212   CeedBasisDestroy(&basis_u);
213   CeedBasisDestroy(&basis_x);
214   CeedVectorDestroy(&x_ref_points);
215   CeedVectorDestroy(&q_data_points);
216   PetscFunctionReturn(PETSC_SUCCESS);
217 }
218 
219 PetscErrorCode DMSwarmCeedContextDestroy(DMSwarmCeedContext *ctx) {
220   PetscFunctionBeginUser;
221   CeedDestroy(&(*ctx)->ceed);
222   CeedVectorDestroy(&(*ctx)->u_mesh);
223   CeedVectorDestroy(&(*ctx)->v_mesh);
224   CeedVectorDestroy(&(*ctx)->u_points);
225   CeedOperatorDestroy(&(*ctx)->op_mesh_to_points);
226   CeedOperatorDestroy(&(*ctx)->op_points_to_mesh);
227   CeedOperatorDestroy(&(*ctx)->op_mass);
228   PetscCall(PetscFree(*ctx));
229   PetscFunctionReturn(PETSC_SUCCESS);
230 }
231 
232 // ------------------------------------------------------------------------------------------------
233 // PETSc-libCEED memory space utilities
234 // ------------------------------------------------------------------------------------------------
235 PetscErrorCode DMSwarmPICFieldP2C(DM dm_swarm, const char *field, CeedVector x_ceed) {
236   PetscScalar *x;
237 
238   PetscFunctionBeginUser;
239   PetscCall(DMSwarmGetField(dm_swarm, field, NULL, NULL, (void **)&x));
240   CeedVectorSetArray(x_ceed, CEED_MEM_HOST, CEED_USE_POINTER, (CeedScalar *)x);
241   PetscFunctionReturn(PETSC_SUCCESS);
242 }
243 
244 PetscErrorCode DMSwarmPICFieldC2P(DM dm_swarm, const char *field, CeedVector x_ceed) {
245   PetscScalar *x;
246 
247   PetscFunctionBeginUser;
248   CeedVectorTakeArray(x_ceed, CEED_MEM_HOST, (CeedScalar **)&x);
249   PetscCall(DMSwarmRestoreField(dm_swarm, field, NULL, NULL, (void **)&x));
250   PetscFunctionReturn(PETSC_SUCCESS);
251 }
252 
253 // ------------------------------------------------------------------------------------------------
254 // Swarm point location utility
255 // ------------------------------------------------------------------------------------------------
256 PetscErrorCode DMSwarmInitalizePointLocations(DM dm_swarm, PointSwarmType point_swarm_type, PetscInt num_points, PetscInt num_points_per_cell) {
257   PetscFunctionBeginUser;
258   switch (point_swarm_type) {
259     case SWARM_GAUSS:
260     case SWARM_UNIFORM: {
261       // -- Set gauss or uniform point locations in each cell
262       PetscInt    num_points_per_cell_1d = round(cbrt(num_points_per_cell * 1.0)), dim = 3;
263       PetscScalar point_coords[num_points_per_cell * 3];
264       CeedScalar  points_1d[num_points_per_cell_1d], weights_1d[num_points_per_cell_1d];
265 
266       if (point_swarm_type == SWARM_GAUSS) {
267         PetscCall(CeedGaussQuadrature(num_points_per_cell_1d, points_1d, weights_1d));
268       } else {
269         for (PetscInt i = 0; i < num_points_per_cell_1d; i++) points_1d[i] = 2.0 * (PetscReal)(i + 1) / (PetscReal)(num_points_per_cell_1d + 1) - 1;
270       }
271       for (PetscInt i = 0; i < num_points_per_cell_1d; i++) {
272         for (PetscInt j = 0; j < num_points_per_cell_1d; j++) {
273           for (PetscInt k = 0; k < num_points_per_cell_1d; k++) {
274             PetscInt p = (i * num_points_per_cell_1d + j) * num_points_per_cell_1d + k;
275 
276             point_coords[p * dim + 0] = points_1d[i];
277             point_coords[p * dim + 1] = points_1d[j];
278             point_coords[p * dim + 2] = points_1d[k];
279           }
280         }
281       }
282       PetscCall(DMSwarmSetPointCoordinatesCellwise(dm_swarm, num_points_per_cell_1d * num_points_per_cell_1d * num_points_per_cell_1d, point_coords));
283     } break;
284     case SWARM_CELL_RANDOM: {
285       // -- Set points randomly in each cell
286       PetscInt     dim = 3, num_cells_total = 1, num_cells[] = {1, 1, 1};
287       PetscScalar *point_coords;
288       PetscRandom  rng;
289 
290       PetscOptionsBegin(PetscObjectComm((PetscObject)dm_swarm), NULL, "libCEED example using PETSc with DMSwarm", NULL);
291 
292       PetscCall(PetscOptionsInt("-dm_plex_dim", "Background mesh dimension", NULL, dim, &dim, NULL));
293       PetscCall(PetscOptionsIntArray("-dm_plex_box_faces", "Number of cells", NULL, num_cells, &dim, NULL));
294 
295       PetscOptionsEnd();
296 
297       PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm_swarm), &rng));
298 
299       num_cells_total = num_cells[0] * num_cells[1] * num_cells[2];
300       PetscCall(DMSwarmGetField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&point_coords));
301       for (PetscInt c = 0; c < num_cells_total; c++) {
302         PetscInt cell_index[3] = {c % num_cells[0], (c / num_cells[0]) % num_cells[1], (c / num_cells[0] / num_cells[1]) % num_cells[2]};
303 
304         for (PetscInt p = 0; p < num_points_per_cell; p++) {
305           PetscInt    point_index = c * num_points_per_cell + p;
306           PetscScalar random_value;
307 
308           for (PetscInt i = 0; i < dim; i++) {
309             PetscCall(PetscRandomGetValue(rng, &random_value));
310             point_coords[point_index * dim + i] = -1.0 + cell_index[i] * 2.0 / (num_cells[i] + 1.0) + random_value;
311           }
312         }
313       }
314       PetscCall(DMSwarmRestoreField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&point_coords));
315       PetscCall(PetscRandomDestroy(&rng));
316     } break;
317     case SWARM_SINUSOIDAL: {
318       // -- Set points distributed per sinusoidal functions
319       PetscInt     dim = 3;
320       PetscScalar *point_coords;
321       DM           dm_mesh;
322 
323       PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
324       PetscCall(DMGetDimension(dm_mesh, &dim));
325 
326       PetscCall(DMSwarmGetField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&point_coords));
327       for (PetscInt p = 0; p < num_points; p++) {
328         point_coords[p * dim + 0] = -PetscCosReal((PetscReal)(p + 1) / (PetscReal)(num_points + 1) * PETSC_PI);
329         if (dim > 1) point_coords[p * dim + 1] = -PetscSinReal((PetscReal)(p + 1) / (PetscReal)(num_points + 1) * PETSC_PI);
330         if (dim > 2) point_coords[p * dim + 2] = PetscSinReal((PetscReal)(p + 1) / (PetscReal)(num_points + 1) * PETSC_PI);
331       }
332       PetscCall(DMSwarmRestoreField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&point_coords));
333     } break;
334   }
335   PetscCall(DMSwarmMigrate(dm_swarm, PETSC_TRUE));
336   PetscFunctionReturn(PETSC_SUCCESS);
337 }
338 
339 /*@C
340   DMSwarmCreateReferenceCoordinates - Compute the cell reference coordinates for local DMSwarm points.
341 
342   Collective
343 
344   Input Parameter:
345 . dm_swarm  - the `DMSwarm`
346 
347   Output Parameters:
348 + is_points    - The IS object for indexing into points per cell
349 - X_points_ref - Vec holding the cell reference coordinates for local DMSwarm points
350 
351 The index set contains ranges of indices for each local cell. This range contains the indices of every point in the cell.
352 
353 ```
354 total_num_cells
355 cell_0_start_index
356 cell_1_start_index
357 cell_2_start_index
358 ...
359 cell_n_start_index
360 cell_n_stop_index
361 cell_0_point_0
362 cell_0_point_0
363 ...
364 cell_n_point_m
365 ```
366 
367   Level: beginner
368 
369 .seealso: `DMSwarm`
370 @*/
371 PetscErrorCode DMSwarmCreateReferenceCoordinates(DM dm_swarm, IS *is_points, Vec *X_points_ref) {
372   PetscInt           cell_start, cell_end, num_cells_local, dim, num_points_local, *cell_points, points_offset;
373   PetscScalar       *coords_points_ref;
374   const PetscScalar *coords_points_true;
375   DM                 dm_mesh;
376 
377   PetscFunctionBeginUser;
378   PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
379 
380   // Create vector to hold reference coordinates
381   {
382     Vec X_points_true;
383 
384     PetscCall(DMSwarmCreateLocalVectorFromField(dm_swarm, DMSwarmPICField_coor, &X_points_true));
385     PetscCall(VecDuplicate(X_points_true, X_points_ref));
386     PetscCall(DMSwarmDestroyLocalVectorFromField(dm_swarm, DMSwarmPICField_coor, &X_points_true));
387   }
388 
389   // Allocate index set array
390   PetscCall(DMPlexGetHeightStratum(dm_mesh, 0, &cell_start, &cell_end));
391   num_cells_local = cell_end - cell_start;
392   points_offset   = num_cells_local + 2;
393   PetscCall(VecGetLocalSize(*X_points_ref, &num_points_local));
394   PetscCall(DMGetDimension(dm_mesh, &dim));
395   num_points_local /= dim;
396   PetscCall(PetscMalloc1(num_points_local + num_cells_local + 2, &cell_points));
397   cell_points[0] = num_cells_local;
398 
399   // Get reference coordinates for each swarm point wrt the elements in the background mesh
400   PetscCall(DMSwarmSortGetAccess(dm_swarm));
401   PetscCall(DMSwarmGetField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&coords_points_true));
402   PetscCall(VecGetArray(*X_points_ref, &coords_points_ref));
403   for (PetscInt cell = cell_start, num_points_processed = 0; cell < cell_end; cell++) {
404     PetscInt *points_in_cell, num_points_in_cell, local_cell = cell - cell_start;
405     PetscReal v[3], J[9], invJ[9], detJ, v0_ref[3] = {-1.0, -1.0, -1.0};
406 
407     PetscCall(DMSwarmSortGetPointsPerCell(dm_swarm, cell, &num_points_in_cell, &points_in_cell));
408     // -- Reference coordinates for swarm points in background mesh element
409     PetscCall(DMPlexComputeCellGeometryFEM(dm_mesh, cell, NULL, v, J, invJ, &detJ));
410     cell_points[local_cell + 1] = num_points_processed + points_offset;
411     for (PetscInt p = 0; p < num_points_in_cell; p++) {
412       const CeedInt point = points_in_cell[p];
413 
414       cell_points[points_offset + (num_points_processed++)] = point;
415       CoordinatesRealToRef(dim, dim, v0_ref, v, invJ, &coords_points_true[point * dim], &coords_points_ref[point * dim]);
416     }
417 
418     // -- Cleanup
419     PetscCall(PetscFree(points_in_cell));
420   }
421   cell_points[points_offset - 1] = num_points_local + points_offset;
422 
423   // Cleanup
424   PetscCall(DMSwarmRestoreField(dm_swarm, DMSwarmPICField_coor, NULL, NULL, (void **)&coords_points_true));
425   PetscCall(VecRestoreArray(*X_points_ref, &coords_points_ref));
426   PetscCall(DMSwarmSortRestoreAccess(dm_swarm));
427 
428   // Create index set
429   PetscCall(ISCreateGeneral(PETSC_COMM_SELF, num_points_local + points_offset, cell_points, PETSC_OWN_POINTER, is_points));
430   PetscFunctionReturn(PETSC_SUCCESS);
431 }
432 
433 // ------------------------------------------------------------------------------------------------
434 // RHS for Swarm to Mesh projection
435 // ------------------------------------------------------------------------------------------------
436 PetscErrorCode DMSwarmCreateProjectionRHS(DM dm_swarm, const char *field, Vec U_points, Vec B_mesh) {
437   PetscMemType       B_mem_type, U_mem_type;
438   DM                 dm_mesh;
439   Vec                B_mesh_loc;
440   PetscBool          has_u_points;
441   DMSwarmCeedContext swarm_ceed_context;
442 
443   PetscFunctionBeginUser;
444   // Get mesh DM
445   PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
446   PetscCall(DMGetApplicationContext(dm_mesh, (void *)&swarm_ceed_context));
447 
448   // Get swarm access
449   has_u_points = !!U_points;
450   if (!has_u_points) {
451     PetscCall(DMSwarmSortGetAccess(dm_swarm));
452     PetscCall(DMSwarmCreateLocalVectorFromField(dm_swarm, field, &U_points));
453   }
454 
455   // Get mesh values
456   PetscCall(VecReadP2C(U_points, &U_mem_type, swarm_ceed_context->u_points));
457   PetscCall(DMGetLocalVector(dm_mesh, &B_mesh_loc));
458   PetscCall(VecZeroEntries(B_mesh_loc));
459   PetscCall(VecP2C(B_mesh_loc, &B_mem_type, swarm_ceed_context->v_mesh));
460 
461   // Interpolate field from swarm points to mesh
462   CeedOperatorApply(swarm_ceed_context->op_points_to_mesh, swarm_ceed_context->u_points, swarm_ceed_context->v_mesh, CEED_REQUEST_IMMEDIATE);
463 
464   // Restore PETSc Vecs and Local to Global
465   PetscCall(VecReadC2P(swarm_ceed_context->u_points, U_mem_type, U_points));
466   PetscCall(VecC2P(swarm_ceed_context->v_mesh, B_mem_type, B_mesh_loc));
467   PetscCall(VecZeroEntries(B_mesh));
468   PetscCall(DMLocalToGlobal(dm_mesh, B_mesh_loc, ADD_VALUES, B_mesh));
469 
470   // Restore swarm access
471   if (!has_u_points) {
472     PetscCall(DMSwarmDestroyLocalVectorFromField(dm_swarm, field, &U_points));
473     PetscCall(DMSwarmSortRestoreAccess(dm_swarm));
474   }
475 
476   // Cleanup
477   PetscCall(DMRestoreLocalVector(dm_mesh, &B_mesh_loc));
478   PetscFunctionReturn(PETSC_SUCCESS);
479 }
480 
481 // ------------------------------------------------------------------------------------------------
482 // Swarm "mass matrix"
483 // ------------------------------------------------------------------------------------------------
484 PetscErrorCode MatMult_SwarmMass(Mat A, Vec U_mesh, Vec V_mesh) {
485   PetscMemType       U_mem_type, V_mem_type;
486   DM                 dm_mesh;
487   Vec                U_mesh_loc, V_mesh_loc;
488   DMSwarmCeedContext swarm_ceed_context;
489 
490   PetscFunctionBeginUser;
491   // Get mesh DM
492   PetscCall(MatGetDM(A, &dm_mesh));
493   PetscCall(DMGetApplicationContext(dm_mesh, (void *)&swarm_ceed_context));
494 
495   // Global to Local and get PETSc Vec access
496   PetscCall(DMGetLocalVector(dm_mesh, &U_mesh_loc));
497   PetscCall(VecZeroEntries(U_mesh_loc));
498   PetscCall(DMGlobalToLocal(dm_mesh, U_mesh, INSERT_VALUES, U_mesh_loc));
499   PetscCall(VecReadP2C(U_mesh_loc, &U_mem_type, swarm_ceed_context->u_mesh));
500   PetscCall(DMGetLocalVector(dm_mesh, &V_mesh_loc));
501   PetscCall(VecZeroEntries(V_mesh_loc));
502   PetscCall(VecP2C(V_mesh_loc, &V_mem_type, swarm_ceed_context->v_mesh));
503 
504   // Apply swarm mass operator
505   CeedOperatorApply(swarm_ceed_context->op_mass, swarm_ceed_context->u_mesh, swarm_ceed_context->v_mesh, CEED_REQUEST_IMMEDIATE);
506 
507   // Restore PETSc Vecs and Local to Global
508   PetscCall(VecReadC2P(swarm_ceed_context->u_mesh, U_mem_type, U_mesh_loc));
509   PetscCall(VecC2P(swarm_ceed_context->v_mesh, V_mem_type, V_mesh_loc));
510   PetscCall(VecZeroEntries(V_mesh));
511   PetscCall(DMLocalToGlobal(dm_mesh, V_mesh_loc, ADD_VALUES, V_mesh));
512 
513   // Cleanup
514   PetscCall(DMRestoreLocalVector(dm_mesh, &U_mesh_loc));
515   PetscCall(DMRestoreLocalVector(dm_mesh, &V_mesh_loc));
516   PetscFunctionReturn(PETSC_SUCCESS);
517 }
518 
519 // ------------------------------------------------------------------------------------------------
520 // Swarm to mesh projection
521 // ------------------------------------------------------------------------------------------------
522 PetscErrorCode DMSwarmProjectFromSwarmToCells(DM dm_swarm, const char *field, Vec U_points, Vec U_mesh) {
523   PetscBool          test_mode;
524   Vec                B_mesh;
525   Mat                M;
526   KSP                ksp;
527   DM                 dm_mesh;
528   DMSwarmCeedContext swarm_ceed_context;
529   MPI_Comm           comm;
530 
531   PetscFunctionBeginUser;
532   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Swarm-to-Mesh Projection Options", NULL);
533   PetscCall(PetscOptionsBool("-test", "Testing mode (do not print unless error is large)", NULL, test_mode, &test_mode, NULL));
534   PetscOptionsEnd();
535 
536   comm = PetscObjectComm((PetscObject)dm_swarm);
537   PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
538   PetscCall(DMGetApplicationContext(dm_mesh, (void *)&swarm_ceed_context));
539   PetscCall(VecDuplicate(U_mesh, &B_mesh));
540 
541   // Setup "mass matrix"
542   {
543     PetscInt l_size, g_size;
544 
545     PetscCall(VecGetLocalSize(U_mesh, &l_size));
546     PetscCall(VecGetSize(U_mesh, &g_size));
547     PetscCall(MatCreateShell(comm, l_size, l_size, g_size, g_size, swarm_ceed_context, &M));
548     PetscCall(MatSetDM(M, dm_mesh));
549     PetscCall(MatShellSetOperation(M, MATOP_MULT, (void (*)(void))MatMult_SwarmMass));
550   }
551 
552   // Setup KSP
553   {
554     PC pc;
555 
556     PetscCall(KSPCreate(comm, &ksp));
557     PetscCall(KSPGetPC(ksp, &pc));
558     PetscCall(PCSetType(pc, PCJACOBI));
559     PetscCall(PCJacobiSetType(pc, PC_JACOBI_ROWSUM));
560     PetscCall(KSPSetType(ksp, KSPCG));
561     PetscCall(KSPSetNormType(ksp, KSP_NORM_NATURAL));
562     PetscCall(KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT));
563     PetscCall(KSPSetOperators(ksp, M, M));
564     PetscCall(KSPSetFromOptions(ksp));
565     PetscCall(PetscObjectSetName((PetscObject)ksp, "Swarm-to-Mesh Projection"));
566     PetscCall(KSPViewFromOptions(ksp, NULL, "-ksp_projection_view"));
567   }
568 
569   // Setup RHS
570   PetscCall(DMSwarmCreateProjectionRHS(dm_swarm, field, U_points, B_mesh));
571 
572   // Solve
573   PetscCall(VecZeroEntries(U_mesh));
574   PetscCall(KSPSolve(ksp, B_mesh, U_mesh));
575 
576   // KSP summary
577   {
578     KSPType            ksp_type;
579     KSPConvergedReason reason;
580     PetscReal          rnorm;
581     PetscInt           its;
582     PetscCall(KSPGetType(ksp, &ksp_type));
583     PetscCall(KSPGetConvergedReason(ksp, &reason));
584     PetscCall(KSPGetIterationNumber(ksp, &its));
585     PetscCall(KSPGetResidualNorm(ksp, &rnorm));
586 
587     if (!test_mode || reason < 0 || rnorm > 1e-8) {
588       PetscCall(PetscPrintf(comm,
589                             "Swarm-to-Mesh Projection KSP Solve:\n"
590                             "  KSP type: %s\n"
591                             "  KSP convergence: %s\n"
592                             "  Total KSP iterations: %" PetscInt_FMT "\n"
593                             "  Final rnorm: %e\n",
594                             ksp_type, KSPConvergedReasons[reason], its, (double)rnorm));
595     }
596   }
597 
598   // Optional viewing
599   PetscCall(KSPViewFromOptions(ksp, NULL, "-ksp_view"));
600 
601   // Cleanup
602   PetscCall(VecDestroy(&B_mesh));
603   PetscCall(MatDestroy(&M));
604   PetscCall(KSPDestroy(&ksp));
605   PetscFunctionReturn(PETSC_SUCCESS);
606 }
607 
608 // ------------------------------------------------------------------------------------------------
609 // BP setup
610 // ------------------------------------------------------------------------------------------------
611 PetscErrorCode SetupProblemSwarm(DM dm_swarm, Ceed ceed, BPData bp_data, CeedData data, PetscBool setup_rhs, Vec rhs, Vec target) {
612   DM                  dm_mesh, dm_coord;
613   CeedElemRestriction elem_restr_u_mesh, elem_restr_x_mesh, elem_restr_x_points, elem_restr_u_points, elem_restr_q_data_points;
614   CeedBasis           basis_u, basis_x;
615   CeedVector          x_coord, x_ref_points, q_data_points;
616   CeedInt             num_comp, q_data_size = bp_data.q_data_size, dim, X_loc_len;
617   CeedScalar          R = 1;                         // radius of the sphere
618   CeedScalar          l = 1.0 / PetscSqrtReal(3.0);  // half edge of the inscribed cube
619   Vec                 X_loc;
620   PetscMemType        X_mem_type;
621 
622   PetscFunctionBeginUser;
623   PetscCall(DMSwarmGetCellDM(dm_swarm, &dm_mesh));
624   PetscCall(DMGetCoordinateDM(dm_mesh, &dm_coord));
625 
626   // Get coordinates
627   PetscCall(DMGetCoordinatesLocal(dm_mesh, &X_loc));
628   PetscCall(VecGetLocalSize(X_loc, &X_loc_len));
629   CeedVectorCreate(ceed, X_loc_len, &x_coord);
630   PetscCall(VecReadP2C(X_loc, &X_mem_type, x_coord));
631 
632   // Background mesh objects
633   PetscCall(CreateBasisFromPlex(ceed, dm_mesh, NULL, 0, 0, 0, bp_data, &basis_u));
634   PetscCall(CreateBasisFromPlex(ceed, dm_coord, NULL, 0, 0, 0, bp_data, &basis_x));
635   PetscCall(CreateRestrictionFromPlex(ceed, dm_mesh, 0, NULL, 0, &elem_restr_u_mesh));
636   PetscCall(CreateRestrictionFromPlex(ceed, dm_coord, 0, NULL, 0, &elem_restr_x_mesh));
637 
638   CeedElemRestrictionCreateVector(elem_restr_u_mesh, &data->x_ceed, NULL);
639   CeedElemRestrictionCreateVector(elem_restr_u_mesh, &data->y_ceed, NULL);
640 
641   // Swarm objects
642   {
643     const PetscInt *cell_points;
644     IS              is_points;
645     Vec             X_ref;
646     CeedInt         num_elem;
647 
648     PetscCall(DMSwarmCreateReferenceCoordinates(dm_swarm, &is_points, &X_ref));
649     PetscCall(DMGetDimension(dm_mesh, &dim));
650     CeedElemRestrictionGetNumElements(elem_restr_u_mesh, &num_elem);
651     CeedElemRestrictionGetNumComponents(elem_restr_u_mesh, &num_comp);
652 
653     PetscCall(ISGetIndices(is_points, &cell_points));
654     PetscInt num_points = cell_points[num_elem + 1] - num_elem - 2;
655     CeedInt  offsets[num_elem + 1 + num_points];
656 
657     for (PetscInt i = 0; i < num_elem + 1; i++) offsets[i] = cell_points[i + 1] - 1;
658     for (PetscInt i = num_elem + 1; i < num_points + num_elem + 1; i++) offsets[i] = cell_points[i + 1];
659     PetscCall(ISRestoreIndices(is_points, &cell_points));
660 
661     // -- Points restrictions
662     CeedElemRestrictionCreateAtPoints(ceed, num_elem, num_points, num_comp, num_points * num_comp, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
663                                       &elem_restr_u_points);
664     CeedElemRestrictionCreateAtPoints(ceed, num_elem, num_points, dim, num_points * dim, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
665                                       &elem_restr_x_points);
666     CeedElemRestrictionCreateAtPoints(ceed, num_elem, num_points, q_data_size, num_points * q_data_size, CEED_MEM_HOST, CEED_COPY_VALUES, offsets,
667                                       &elem_restr_q_data_points);
668 
669     // -- Points vectors
670     CeedElemRestrictionCreateVector(elem_restr_q_data_points, &q_data_points, NULL);
671 
672     // -- Ref coordinates
673     {
674       PetscMemType       X_mem_type;
675       const PetscScalar *x;
676 
677       CeedVectorCreate(ceed, num_points * dim, &x_ref_points);
678 
679       PetscCall(VecGetArrayReadAndMemType(X_ref, (const PetscScalar **)&x, &X_mem_type));
680       CeedVectorSetArray(x_ref_points, MemTypeP2C(X_mem_type), CEED_COPY_VALUES, (CeedScalar *)x);
681       PetscCall(VecRestoreArrayReadAndMemType(X_ref, (const PetscScalar **)&x));
682     }
683 
684     // Create Q data
685     {
686       CeedQFunction qf_setup;
687       CeedOperator  op_setup;
688 
689       // Setup geometric scaling
690       CeedQFunctionCreateInterior(ceed, 1, bp_data.setup_geo, bp_data.setup_geo_loc, &qf_setup);
691       CeedQFunctionAddInput(qf_setup, "x", dim, CEED_EVAL_INTERP);
692       CeedQFunctionAddInput(qf_setup, "dx", dim * dim, CEED_EVAL_GRAD);
693       CeedQFunctionAddInput(qf_setup, "weight", 1, CEED_EVAL_WEIGHT);
694       CeedQFunctionAddOutput(qf_setup, "qdata", q_data_size, CEED_EVAL_NONE);
695 
696       CeedOperatorCreateAtPoints(ceed, qf_setup, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE, &op_setup);
697       CeedOperatorSetField(op_setup, "x", elem_restr_x_mesh, basis_x, CEED_VECTOR_ACTIVE);
698       CeedOperatorSetField(op_setup, "dx", elem_restr_x_mesh, basis_x, CEED_VECTOR_ACTIVE);
699       CeedOperatorSetField(op_setup, "weight", CEED_ELEMRESTRICTION_NONE, basis_x, CEED_VECTOR_NONE);
700       CeedOperatorSetField(op_setup, "qdata", elem_restr_q_data_points, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE);
701       CeedOperatorAtPointsSetPoints(op_setup, elem_restr_x_points, x_ref_points);
702 
703       CeedOperatorApply(op_setup, x_coord, q_data_points, CEED_REQUEST_IMMEDIATE);
704 
705       // Cleanup
706       CeedQFunctionDestroy(&qf_setup);
707       CeedOperatorDestroy(&op_setup);
708     }
709 
710     // Cleanup
711     PetscCall(ISDestroy(&is_points));
712     PetscCall(VecDestroy(&X_ref));
713   }
714 
715   // Set up PDE operator
716 
717   CeedQFunction qf_apply;
718   CeedOperator  op_apply;
719   CeedInt       in_scale  = bp_data.in_mode == CEED_EVAL_GRAD ? dim : 1;
720   CeedInt       out_scale = bp_data.out_mode == CEED_EVAL_GRAD ? dim : 1;
721 
722   CeedQFunctionCreateInterior(ceed, 1, bp_data.apply, bp_data.apply_loc, &qf_apply);
723   CeedQFunctionAddInput(qf_apply, "u", num_comp * in_scale, bp_data.in_mode);
724   CeedQFunctionAddInput(qf_apply, "qdata", q_data_size, CEED_EVAL_NONE);
725   CeedQFunctionAddOutput(qf_apply, "v", num_comp * out_scale, bp_data.out_mode);
726 
727   // Create the mass or diff operator
728   CeedOperatorCreateAtPoints(ceed, qf_apply, NULL, NULL, &op_apply);
729   CeedOperatorSetField(op_apply, "u", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
730   CeedOperatorSetField(op_apply, "qdata", elem_restr_q_data_points, CEED_BASIS_NONE, q_data_points);
731   CeedOperatorSetField(op_apply, "v", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
732   CeedOperatorAtPointsSetPoints(op_apply, elem_restr_x_points, x_ref_points);
733 
734   // Set up RHS if needed
735   if (setup_rhs) {
736     CeedQFunction qf_setup_rhs;
737     CeedOperator  op_setup_rhs;
738     Vec           rhs_loc;
739     CeedVector    rhs_ceed, target_ceed;
740     PetscMemType  rhs_mem_type, target_mem_type;
741 
742     // Create RHS vector
743     PetscCall(DMCreateLocalVector(dm_mesh, &rhs_loc));
744 
745     CeedElemRestrictionCreateVector(elem_restr_u_points, &target_ceed, NULL);
746     CeedElemRestrictionCreateVector(elem_restr_u_mesh, &rhs_ceed, NULL);
747 
748     // Create the q-function that sets up the RHS and true solution
749     CeedQFunctionCreateInterior(ceed, 1, bp_data.setup_rhs, bp_data.setup_rhs_loc, &qf_setup_rhs);
750     CeedQFunctionAddInput(qf_setup_rhs, "x", dim, CEED_EVAL_INTERP);
751     CeedQFunctionAddInput(qf_setup_rhs, "qdata", q_data_size, CEED_EVAL_NONE);
752     CeedQFunctionAddOutput(qf_setup_rhs, "true solution", num_comp, CEED_EVAL_NONE);
753     CeedQFunctionAddOutput(qf_setup_rhs, "rhs", num_comp, CEED_EVAL_INTERP);
754 
755     // Create the operator that builds the RHS and true solution
756     CeedOperatorCreateAtPoints(ceed, qf_setup_rhs, NULL, NULL, &op_setup_rhs);
757     CeedOperatorSetField(op_setup_rhs, "x", elem_restr_x_mesh, basis_x, CEED_VECTOR_ACTIVE);
758     CeedOperatorSetField(op_setup_rhs, "qdata", elem_restr_q_data_points, CEED_BASIS_NONE, q_data_points);
759     CeedOperatorSetField(op_setup_rhs, "true solution", elem_restr_u_points, CEED_BASIS_NONE, target_ceed);
760     CeedOperatorSetField(op_setup_rhs, "rhs", elem_restr_u_mesh, basis_u, CEED_VECTOR_ACTIVE);
761     CeedOperatorAtPointsSetPoints(op_setup_rhs, elem_restr_x_points, x_ref_points);
762 
763     // Set up the libCEED context
764     CeedQFunctionContext ctx_rhs_setup;
765     CeedQFunctionContextCreate(ceed, &ctx_rhs_setup);
766     CeedScalar rhs_setup_data[2] = {R, l};
767     CeedQFunctionContextSetData(ctx_rhs_setup, CEED_MEM_HOST, CEED_COPY_VALUES, sizeof rhs_setup_data, &rhs_setup_data);
768     CeedQFunctionSetContext(qf_setup_rhs, ctx_rhs_setup);
769     CeedQFunctionContextDestroy(&ctx_rhs_setup);
770 
771     // Setup RHS and target
772     PetscCall(VecP2C(rhs_loc, &rhs_mem_type, rhs_ceed));
773     PetscCall(VecP2C(target, &target_mem_type, target_ceed));
774     CeedOperatorApply(op_setup_rhs, x_coord, rhs_ceed, CEED_REQUEST_IMMEDIATE);
775     PetscCall(VecC2P(rhs_ceed, rhs_mem_type, rhs_loc));
776     PetscCall(VecC2P(target_ceed, target_mem_type, target));
777 
778     // Local-to-global
779     PetscCall(VecZeroEntries(rhs));
780     PetscCall(DMLocalToGlobal(dm_mesh, rhs_loc, ADD_VALUES, rhs));
781 
782     PetscCall(VecViewFromOptions(rhs, NULL, "-rhs_view"));
783 
784     // Cleanup
785     PetscCall(DMRestoreLocalVector(dm_mesh, &rhs_loc));
786     CeedVectorDestroy(&rhs_ceed);
787     CeedVectorDestroy(&target_ceed);
788     CeedQFunctionDestroy(&qf_setup_rhs);
789     CeedOperatorDestroy(&op_setup_rhs);
790   }
791 
792   // Save libCEED data
793   data->basis_x         = basis_x;
794   data->basis_u         = basis_u;
795   data->elem_restr_x    = elem_restr_x_mesh;
796   data->elem_restr_u    = elem_restr_u_mesh;
797   data->elem_restr_u_i  = elem_restr_u_points;
798   data->elem_restr_qd_i = elem_restr_q_data_points;
799   data->qf_apply        = qf_apply;
800   data->op_apply        = op_apply;
801   data->q_data          = q_data_points;
802   data->q_data_size     = q_data_size;
803 
804   // Cleanup
805   PetscCall(VecReadC2P(x_coord, X_mem_type, X_loc));
806   CeedVectorDestroy(&x_coord);
807   CeedVectorDestroy(&x_ref_points);
808   CeedElemRestrictionDestroy(&elem_restr_x_points);
809 
810   PetscFunctionReturn(PETSC_SUCCESS);
811 }
812