xref: /libCEED/examples/petsc/src/libceedsetup.c (revision 93b6d8191bb649fce95198206c6bff32567615d1)
1 #include <stdio.h>
2 #include "../include/libceedsetup.h"
3 #include "../include/petscutils.h"
4 
5 // -----------------------------------------------------------------------------
6 // Destroy libCEED operator objects
7 // -----------------------------------------------------------------------------
8 PetscErrorCode CeedDataDestroy(CeedInt i, CeedData data) {
9   int ierr;
10 
11   PetscFunctionBeginUser;
12   CeedVectorDestroy(&data->q_data);
13   CeedVectorDestroy(&data->x_ceed);
14   CeedVectorDestroy(&data->y_ceed);
15   CeedBasisDestroy(&data->basis_x);
16   CeedBasisDestroy(&data->basis_u);
17   CeedElemRestrictionDestroy(&data->elem_restr_u);
18   CeedElemRestrictionDestroy(&data->elem_restr_x);
19   CeedElemRestrictionDestroy(&data->elem_restr_u_i);
20   CeedElemRestrictionDestroy(&data->elem_restr_qd_i);
21   CeedQFunctionDestroy(&data->qf_apply);
22   CeedOperatorDestroy(&data->op_apply);
23   if (i > 0) {
24     CeedOperatorDestroy(&data->op_prolong);
25     CeedBasisDestroy(&data->basis_c_to_f);
26     CeedOperatorDestroy(&data->op_restrict);
27   }
28   ierr = PetscFree(data); CHKERRQ(ierr);
29 
30   PetscFunctionReturn(0);
31 };
32 
33 // -----------------------------------------------------------------------------
34 // Set up libCEED for a given degree
35 // -----------------------------------------------------------------------------
36 PetscErrorCode SetupLibceedByDegree(DM dm, Ceed ceed, CeedInt degree,
37                                     CeedInt topo_dim, CeedInt q_extra,
38                                     PetscInt num_comp_x, PetscInt num_comp_u,
39                                     PetscInt g_size, PetscInt xl_size,
40                                     BPData bp_data, CeedData data,
41                                     PetscBool setup_rhs, CeedVector rhs_ceed,
42                                     CeedVector *target) {
43   int ierr;
44   DM dm_coord;
45   Vec coords;
46   const PetscScalar *coord_array;
47   CeedBasis basis_x, basis_u;
48   CeedElemRestriction elem_restr_x, elem_restr_u, elem_restr_u_i, elem_restr_qd_i;
49   CeedQFunction qf_setup_geo, qf_apply;
50   CeedOperator op_setup_geo, op_apply;
51   CeedVector x_coord, q_data, x_ceed, y_ceed;
52   CeedInt P, Q, num_qpts, c_start, c_end, num_elem,
53           q_data_size = bp_data.q_data_size;
54   CeedScalar R = 1,                      // radius of the sphere
55              l = 1.0/PetscSqrtReal(3.0); // half edge of the inscribed cube
56 
57   PetscFunctionBeginUser;
58   // CEED bases
59   P = degree + 1;
60   Q = P + q_extra;
61   CeedBasisCreateTensorH1Lagrange(ceed, topo_dim, num_comp_u, P, Q,
62                                   bp_data.q_mode,
63                                   &basis_u);
64   CeedBasisCreateTensorH1Lagrange(ceed, topo_dim, num_comp_x, 2, Q,
65                                   bp_data.q_mode,
66                                   &basis_x);
67   CeedBasisGetNumQuadraturePoints(basis_u, &num_qpts);
68 
69   // CEED restrictions
70   ierr = DMSetCoordinateDim(dm, topo_dim); CHKERRQ(ierr);
71   ierr = DMGetCoordinateDM(dm, &dm_coord); CHKERRQ(ierr);
72   ierr = DMPlexSetClosurePermutationTensor(dm_coord, PETSC_DETERMINE, NULL);
73   CHKERRQ(ierr);
74   ierr = CreateRestrictionFromPlex(ceed, dm_coord, 0, 0, 0, &elem_restr_x);
75   CHKERRQ(ierr);
76   ierr = CreateRestrictionFromPlex(ceed, dm, 0, 0, 0, &elem_restr_u);
77   CHKERRQ(ierr);
78 
79   ierr = DMPlexGetHeightStratum(dm, 0, &c_start, &c_end); CHKERRQ(ierr);
80   num_elem = c_end - c_start;
81 
82   CeedElemRestrictionCreateStrided(ceed, num_elem, num_qpts, num_comp_u,
83                                    num_comp_u*num_elem*num_qpts,
84                                    CEED_STRIDES_BACKEND, &elem_restr_u_i);
85   CeedElemRestrictionCreateStrided(ceed, num_elem, num_qpts, q_data_size,
86                                    q_data_size*num_elem*num_qpts,
87                                    CEED_STRIDES_BACKEND, &elem_restr_qd_i);
88 
89   // Element coordinates
90   ierr = DMGetCoordinatesLocal(dm, &coords); CHKERRQ(ierr);
91   ierr = VecGetArrayRead(coords, &coord_array); CHKERRQ(ierr);
92 
93   CeedElemRestrictionCreateVector(elem_restr_x, &x_coord, NULL);
94   CeedVectorSetArray(x_coord, CEED_MEM_HOST, CEED_COPY_VALUES,
95                      (PetscScalar *)coord_array);
96   ierr = VecRestoreArrayRead(coords, &coord_array);
97 
98   // Create the persistent vectors that will be needed in setup and apply
99   CeedVectorCreate(ceed, q_data_size*num_elem*num_qpts, &q_data);
100   CeedVectorCreate(ceed, xl_size, &x_ceed);
101   CeedVectorCreate(ceed, xl_size, &y_ceed);
102 
103   // Create the QFunction that builds the context data
104   CeedQFunctionCreateInterior(ceed, 1, bp_data.setup_geo, bp_data.setup_geo_loc,
105                               &qf_setup_geo);
106   CeedQFunctionAddInput(qf_setup_geo, "x", num_comp_x, CEED_EVAL_INTERP);
107   CeedQFunctionAddInput(qf_setup_geo, "dx", num_comp_x*topo_dim, CEED_EVAL_GRAD);
108   CeedQFunctionAddInput(qf_setup_geo, "weight", 1, CEED_EVAL_WEIGHT);
109   CeedQFunctionAddOutput(qf_setup_geo, "qdata", q_data_size, CEED_EVAL_NONE);
110 
111   // Create the operator that builds the quadrature data
112   CeedOperatorCreate(ceed, qf_setup_geo, NULL, NULL, &op_setup_geo);
113   CeedOperatorSetField(op_setup_geo, "x", elem_restr_x, basis_x,
114                        CEED_VECTOR_ACTIVE);
115   CeedOperatorSetField(op_setup_geo, "dx", elem_restr_x, basis_x,
116                        CEED_VECTOR_ACTIVE);
117   CeedOperatorSetField(op_setup_geo, "weight", CEED_ELEMRESTRICTION_NONE, basis_x,
118                        CEED_VECTOR_NONE);
119   CeedOperatorSetField(op_setup_geo, "qdata", elem_restr_qd_i,
120                        CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
121 
122   // Setup q_data
123   CeedOperatorApply(op_setup_geo, x_coord, q_data, CEED_REQUEST_IMMEDIATE);
124 
125   // Set up PDE operator
126   CeedInt in_scale = bp_data.in_mode == CEED_EVAL_GRAD ? topo_dim : 1;
127   CeedInt out_scale = bp_data.out_mode == CEED_EVAL_GRAD ? topo_dim : 1;
128   CeedQFunctionCreateInterior(ceed, 1, bp_data.apply, bp_data.apply_loc,
129                               &qf_apply);
130   CeedQFunctionAddInput(qf_apply, "u", num_comp_u*in_scale, bp_data.in_mode);
131   CeedQFunctionAddInput(qf_apply, "qdata", q_data_size, CEED_EVAL_NONE);
132   CeedQFunctionAddOutput(qf_apply, "v", num_comp_u*out_scale, bp_data.out_mode);
133 
134   // Create the mass or diff operator
135   CeedOperatorCreate(ceed, qf_apply, NULL, NULL, &op_apply);
136   CeedOperatorSetField(op_apply, "u", elem_restr_u, basis_u, CEED_VECTOR_ACTIVE);
137   CeedOperatorSetField(op_apply, "qdata", elem_restr_qd_i, CEED_BASIS_COLLOCATED,
138                        q_data);
139   CeedOperatorSetField(op_apply, "v", elem_restr_u, basis_u, CEED_VECTOR_ACTIVE);
140 
141   // Set up RHS if needed
142   if (setup_rhs) {
143     CeedQFunction qf_setup_rhs;
144     CeedOperator op_setup_rhs;
145     CeedVectorCreate(ceed, num_elem*num_qpts*num_comp_u, target);
146 
147     // Create the q-function that sets up the RHS and true solution
148     CeedQFunctionCreateInterior(ceed, 1, bp_data.setup_rhs, bp_data.setup_rhs_loc,
149                                 &qf_setup_rhs);
150     CeedQFunctionAddInput(qf_setup_rhs, "x", num_comp_x, CEED_EVAL_INTERP);
151     CeedQFunctionAddInput(qf_setup_rhs, "qdata", q_data_size, CEED_EVAL_NONE);
152     CeedQFunctionAddOutput(qf_setup_rhs, "true solution", num_comp_u,
153                            CEED_EVAL_NONE);
154     CeedQFunctionAddOutput(qf_setup_rhs, "rhs", num_comp_u, CEED_EVAL_INTERP);
155 
156     // Create the operator that builds the RHS and true solution
157     CeedOperatorCreate(ceed, qf_setup_rhs, NULL, NULL, &op_setup_rhs);
158     CeedOperatorSetField(op_setup_rhs, "x", elem_restr_x, basis_x,
159                          CEED_VECTOR_ACTIVE);
160     CeedOperatorSetField(op_setup_rhs, "qdata", elem_restr_qd_i,
161                          CEED_BASIS_COLLOCATED, q_data);
162     CeedOperatorSetField(op_setup_rhs, "true solution", elem_restr_u_i,
163                          CEED_BASIS_COLLOCATED, *target);
164     CeedOperatorSetField(op_setup_rhs, "rhs", elem_restr_u, basis_u,
165                          CEED_VECTOR_ACTIVE);
166 
167     // Set up the libCEED context
168     CeedQFunctionContext ctx_rhs_setup;
169     CeedQFunctionContextCreate(ceed, &ctx_rhs_setup);
170     CeedScalar rhs_setup_data[2] = {R, l};
171     CeedQFunctionContextSetData(ctx_rhs_setup, CEED_MEM_HOST, CEED_COPY_VALUES,
172                                 sizeof rhs_setup_data, &rhs_setup_data);
173     CeedQFunctionSetContext(qf_setup_rhs, ctx_rhs_setup);
174     CeedQFunctionContextDestroy(&ctx_rhs_setup);
175 
176     // Setup RHS and target
177     CeedOperatorApply(op_setup_rhs, x_coord, rhs_ceed, CEED_REQUEST_IMMEDIATE);
178 
179     // Cleanup
180     CeedQFunctionDestroy(&qf_setup_rhs);
181     CeedOperatorDestroy(&op_setup_rhs);
182   }
183 
184   // Cleanup
185   CeedQFunctionDestroy(&qf_setup_geo);
186   CeedOperatorDestroy(&op_setup_geo);
187   CeedVectorDestroy(&x_coord);
188 
189   // Save libCEED data required for level
190   data->basis_x = basis_x; data->basis_u = basis_u;
191   data->elem_restr_x = elem_restr_x;
192   data->elem_restr_u = elem_restr_u;
193   data->elem_restr_u_i = elem_restr_u_i;
194   data->elem_restr_qd_i = elem_restr_qd_i;
195   data->qf_apply = qf_apply;
196   data->op_apply = op_apply;
197   data->q_data = q_data;
198   data->x_ceed = x_ceed;
199   data->y_ceed = y_ceed;
200 
201   PetscFunctionReturn(0);
202 };
203 
204 // -----------------------------------------------------------------------------
205 // Setup libCEED level transfer operator objects
206 // -----------------------------------------------------------------------------
207 PetscErrorCode CeedLevelTransferSetup(Ceed ceed, CeedInt num_levels,
208                                       CeedInt num_comp_u, CeedData *data,
209                                       CeedInt *level_degrees,
210                                       CeedQFunction qf_restrict, CeedQFunction qf_prolong) {
211   PetscFunctionBeginUser;
212   // Return early if num_levels=1
213   if (num_levels == 1)
214     PetscFunctionReturn(0);
215 
216   // Set up each level
217   for (CeedInt i=1; i<num_levels; i++) {
218     // P coarse and P fine
219     CeedInt Pc = level_degrees[i-1] + 1;
220     CeedInt Pf = level_degrees[i] + 1;
221 
222     // Restriction - Fine to corse
223     CeedBasis basis_c_to_f;
224     CeedOperator op_restrict;
225 
226     // Basis
227     CeedBasisCreateTensorH1Lagrange(ceed, 3, num_comp_u, Pc, Pf,
228                                     CEED_GAUSS_LOBATTO, &basis_c_to_f);
229 
230     // Create the restriction operator
231     CeedOperatorCreate(ceed, qf_restrict, CEED_QFUNCTION_NONE,
232                        CEED_QFUNCTION_NONE, &op_restrict);
233     CeedOperatorSetField(op_restrict, "input", data[i]->elem_restr_u,
234                          CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
235     CeedOperatorSetField(op_restrict, "output", data[i-1]->elem_restr_u,
236                          basis_c_to_f, CEED_VECTOR_ACTIVE);
237 
238     // Save libCEED data required for level
239     data[i]->basis_c_to_f = basis_c_to_f;
240     data[i]->op_restrict = op_restrict;
241 
242     // Interpolation - Corse to fine
243     CeedOperator op_prolong;
244 
245     // Create the prolongation operator
246     CeedOperatorCreate(ceed, qf_prolong, CEED_QFUNCTION_NONE,
247                        CEED_QFUNCTION_NONE, &op_prolong);
248     CeedOperatorSetField(op_prolong, "input", data[i-1]->elem_restr_u,
249                          basis_c_to_f, CEED_VECTOR_ACTIVE);
250     CeedOperatorSetField(op_prolong, "output", data[i]->elem_restr_u,
251                          CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
252 
253     // Save libCEED data required for level
254     data[i]->op_prolong = op_prolong;
255   }
256 
257   PetscFunctionReturn(0);
258 };
259 
260 // -----------------------------------------------------------------------------
261