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