xref: /libCEED/examples/fluids/src/turb_spanstats.c (revision 8ed527305783a8e6c48f5ac92760da88017905af)
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 /// @file
8 /// Functions for setting up and performing statistics collection
9 
10 #include "../qfunctions/turb_spanstats.h"
11 
12 #include <petscsf.h>
13 
14 #include "../include/matops.h"
15 #include "../navierstokes.h"
16 #include "ceed/ceed.h"
17 #include "petscerror.h"
18 #include "petsclog.h"
19 #include "petscmat.h"
20 #include "petscsys.h"
21 #include "petscvec.h"
22 #include "petscviewer.h"
23 
24 PetscErrorCode CreateStatsDM(User user, ProblemData *problem, PetscInt degree, SimpleBC bc) {
25   user->spanstats.num_comp_stats = 22;
26   PetscReal domain_min[3], domain_max[3];
27   PetscFunctionBeginUser;
28 
29   // Get spanwise length
30   PetscCall(DMGetBoundingBox(user->dm, domain_min, domain_max));
31   user->spanstats.span_width = domain_max[2] - domain_min[1];
32 
33   // Get DM from surface
34   {
35     DMLabel label;
36     PetscCall(DMGetLabel(user->dm, "Face Sets", &label));
37     PetscCall(DMPlexLabelComplete(user->dm, label));
38     PetscCall(DMPlexFilter(user->dm, label, 1, &user->spanstats.dm));
39     PetscCall(DMProjectCoordinates(user->spanstats.dm, NULL));  // Ensure that a coordinate FE exists
40   }
41 
42   PetscCall(PetscObjectSetName((PetscObject)user->spanstats.dm, "Spanwise_Stats"));
43   PetscCall(DMSetOptionsPrefix(user->spanstats.dm, "spanstats_"));
44   PetscCall(DMSetFromOptions(user->spanstats.dm));
45   PetscCall(DMViewFromOptions(user->spanstats.dm, NULL, "-dm_view"));  // -spanstats_dm_view
46   {
47     PetscFE fe;
48     DMLabel label;
49 
50     PetscCall(PetscFECreateLagrange(PETSC_COMM_SELF, problem->dim - 1, user->spanstats.num_comp_stats, PETSC_FALSE, degree, PETSC_DECIDE, &fe));
51     PetscCall(PetscObjectSetName((PetscObject)fe, "stats"));
52     PetscCall(DMAddField(user->spanstats.dm, NULL, (PetscObject)fe));
53     PetscCall(DMCreateDS(user->spanstats.dm));
54     PetscCall(DMGetLabel(user->spanstats.dm, "Face Sets", &label));
55 
56     PetscCall(DMPlexSetClosurePermutationTensor(user->spanstats.dm, PETSC_DETERMINE, NULL));
57     PetscCall(PetscFEDestroy(&fe));
58   }
59 
60   PetscSection section;
61   PetscCall(DMGetLocalSection(user->spanstats.dm, &section));
62   PetscCall(PetscSectionSetFieldName(section, 0, ""));
63   PetscCall(PetscSectionSetComponentName(section, 0, 0, "Mean Density"));
64   PetscCall(PetscSectionSetComponentName(section, 0, 1, "Mean Pressure"));
65   PetscCall(PetscSectionSetComponentName(section, 0, 2, "Mean Pressure Squared"));
66   PetscCall(PetscSectionSetComponentName(section, 0, 3, "Mean Pressure Velocity X"));
67   PetscCall(PetscSectionSetComponentName(section, 0, 4, "Mean Pressure Velocity Y"));
68   PetscCall(PetscSectionSetComponentName(section, 0, 5, "Mean Pressure Velocity Z"));
69   PetscCall(PetscSectionSetComponentName(section, 0, 6, "Mean Density Temperature"));
70   PetscCall(PetscSectionSetComponentName(section, 0, 7, "Mean Density Temperature Flux X"));
71   PetscCall(PetscSectionSetComponentName(section, 0, 8, "Mean Density Temperature Flux Y"));
72   PetscCall(PetscSectionSetComponentName(section, 0, 9, "Mean Density Temperature Flux Z"));
73   PetscCall(PetscSectionSetComponentName(section, 0, 10, "Mean Momentum X"));
74   PetscCall(PetscSectionSetComponentName(section, 0, 11, "Mean Momentum Y"));
75   PetscCall(PetscSectionSetComponentName(section, 0, 12, "Mean Momentum Z"));
76   PetscCall(PetscSectionSetComponentName(section, 0, 13, "Mean Momentum Flux XX"));
77   PetscCall(PetscSectionSetComponentName(section, 0, 14, "Mean Momentum Flux YY"));
78   PetscCall(PetscSectionSetComponentName(section, 0, 15, "Mean Momentum Flux ZZ"));
79   PetscCall(PetscSectionSetComponentName(section, 0, 16, "Mean Momentum Flux YZ"));
80   PetscCall(PetscSectionSetComponentName(section, 0, 17, "Mean Momentum Flux XZ"));
81   PetscCall(PetscSectionSetComponentName(section, 0, 18, "Mean Momentum Flux XY"));
82   PetscCall(PetscSectionSetComponentName(section, 0, 19, "Mean Velocity X"));
83   PetscCall(PetscSectionSetComponentName(section, 0, 20, "Mean Velocity Y"));
84   PetscCall(PetscSectionSetComponentName(section, 0, 21, "Mean Velocity Z"));
85 
86   PetscFunctionReturn(0);
87 }
88 
89 // Create CeedElemRestriction for collocated data based on associated CeedBasis and CeedElemRestriction
90 // Number of quadrature points is used from the CeedBasis, and number of elements is used from the CeedElemRestriction
91 PetscErrorCode CreateElemRestrColloc(Ceed ceed, CeedInt num_comp, CeedBasis basis, CeedElemRestriction elem_restr_base,
92                                      CeedElemRestriction *elem_restr_collocated, CeedVector *l_vec, CeedVector *e_vec) {
93   CeedInt num_elem_qpts, loc_num_elem;
94   PetscFunctionBeginUser;
95 
96   CeedBasisGetNumQuadraturePoints(basis, &num_elem_qpts);
97   CeedElemRestrictionGetNumElements(elem_restr_base, &loc_num_elem);
98 
99   const CeedInt strides[] = {num_comp, 1, num_elem_qpts * num_comp};
100   CeedElemRestrictionCreateStrided(ceed, loc_num_elem, num_elem_qpts, num_comp, num_comp * loc_num_elem * num_elem_qpts, strides,
101                                    elem_restr_collocated);
102   CeedElemRestrictionCreateVector(*elem_restr_collocated, l_vec, e_vec);
103   PetscFunctionReturn(0);
104 }
105 
106 // Get coordinates of quadrature points
107 PetscErrorCode GetQuadratureCoords(Ceed ceed, DM dm, CeedElemRestriction elem_restr_x, CeedBasis basis_x, CeedVector x_coords, CeedVector *qx_coords,
108                                    PetscInt *total_nqpnts) {
109   CeedQFunction       qf_quad_coords;
110   CeedOperator        op_quad_coords;
111   PetscInt            num_comp_x, loc_num_elem, num_elem_qpts;
112   CeedElemRestriction elem_restr_qx;
113   PetscFunctionBeginUser;
114 
115   // Create Element Restriction and CeedVector for quadrature coordinates
116   CeedBasisGetNumQuadraturePoints(basis_x, &num_elem_qpts);
117   CeedElemRestrictionGetNumElements(elem_restr_x, &loc_num_elem);
118   CeedElemRestrictionGetNumComponents(elem_restr_x, &num_comp_x);
119   *total_nqpnts = num_elem_qpts * loc_num_elem;
120   PetscCall(CreateElemRestrColloc(ceed, num_comp_x, basis_x, elem_restr_x, &elem_restr_qx, qx_coords, NULL));
121 
122   // Create QFunction
123   CeedQFunctionCreateIdentity(ceed, num_comp_x, CEED_EVAL_INTERP, CEED_EVAL_NONE, &qf_quad_coords);
124 
125   // Create Operator
126   CeedOperatorCreate(ceed, qf_quad_coords, NULL, NULL, &op_quad_coords);
127   CeedOperatorSetField(op_quad_coords, "input", elem_restr_x, basis_x, CEED_VECTOR_ACTIVE);
128   CeedOperatorSetField(op_quad_coords, "output", elem_restr_qx, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
129 
130   CeedOperatorApply(op_quad_coords, x_coords, *qx_coords, CEED_REQUEST_IMMEDIATE);
131 
132   CeedQFunctionDestroy(&qf_quad_coords);
133   CeedOperatorDestroy(&op_quad_coords);
134   PetscFunctionReturn(0);
135 }
136 
137 // Create PetscSF for child-to-parent communication
138 PetscErrorCode CreateStatsSF(Ceed ceed, CeedData ceed_data, DM parentdm, DM childdm, PetscSF statssf) {
139   PetscInt   child_num_qpnts, parent_num_qpnts, num_comp_x;
140   CeedVector child_qx_coords, parent_qx_coords;
141   PetscReal *child_coords, *parent_coords;
142   PetscFunctionBeginUser;
143 
144   // Assume that child and parent have the same number of components
145   CeedBasisGetNumComponents(ceed_data->basis_x, &num_comp_x);
146   const PetscInt num_comp_sf = num_comp_x - 1;  // Number of coord components used in the creation of the SF
147 
148   // Get quad_coords for child DM
149   PetscCall(GetQuadratureCoords(ceed, childdm, ceed_data->elem_restr_x, ceed_data->basis_x, ceed_data->x_coord, &child_qx_coords, &child_num_qpnts));
150 
151   // Get quad_coords for parent DM
152   PetscCall(GetQuadratureCoords(ceed, parentdm, ceed_data->spanstats.elem_restr_parent_x, ceed_data->spanstats.basis_x, ceed_data->spanstats.x_coord,
153                                 &parent_qx_coords, &parent_num_qpnts));
154 
155   // Remove z component of coordinates for matching
156   {
157     const PetscReal *child_quad_coords, *parent_quad_coords;
158 
159     CeedVectorGetArrayRead(child_qx_coords, CEED_MEM_HOST, &child_quad_coords);
160     CeedVectorGetArrayRead(parent_qx_coords, CEED_MEM_HOST, &parent_quad_coords);
161 
162     PetscCall(PetscMalloc2(child_num_qpnts * 2, &child_coords, parent_num_qpnts * 2, &parent_coords));
163     for (int i = 0; i < child_num_qpnts; i++) {
164       child_coords[0 + i * num_comp_sf] = child_quad_coords[0 + i * num_comp_x];
165       child_coords[1 + i * num_comp_sf] = child_quad_coords[1 + i * num_comp_x];
166     }
167     for (int i = 0; i < parent_num_qpnts; i++) {
168       parent_coords[0 + i * num_comp_sf] = parent_quad_coords[0 + i * num_comp_x];
169       parent_coords[1 + i * num_comp_sf] = parent_quad_coords[1 + i * num_comp_x];
170     }
171     CeedVectorRestoreArrayRead(child_qx_coords, &child_quad_coords);
172     CeedVectorRestoreArrayRead(parent_qx_coords, &parent_quad_coords);
173   }
174 
175   PetscCall(PetscSFSetGraphFromCoordinates(statssf, parent_num_qpnts, child_num_qpnts, num_comp_sf, 1e-12, parent_coords, child_coords));
176 
177   PetscCall(PetscSFViewFromOptions(statssf, NULL, "-spanstats_sf_view"));
178 
179   PetscCall(PetscFree2(child_coords, parent_coords));
180   CeedVectorDestroy(&child_qx_coords);
181   CeedVectorDestroy(&parent_qx_coords);
182   PetscFunctionReturn(0);
183 }
184 
185 // Compute mass matrix for statistics projection
186 PetscErrorCode SetupL2ProjectionStats(Ceed ceed, User user, CeedData ceed_data) {
187   CeedQFunction qf_mass;
188   CeedOperator  op_mass;
189   CeedInt       num_comp_q, q_data_size;
190   PetscFunctionBeginUser;
191 
192   // CEED Restriction
193   CeedElemRestrictionGetNumComponents(ceed_data->spanstats.elem_restr_parent_stats, &num_comp_q);
194   CeedElemRestrictionGetNumComponents(ceed_data->spanstats.elem_restr_parent_qd, &q_data_size);
195 
196   // Create Mass CeedOperator
197   PetscCall(CreateMassQFunction(ceed, num_comp_q, q_data_size, &qf_mass));
198   CeedOperatorCreate(ceed, qf_mass, NULL, NULL, &op_mass);
199   CeedOperatorSetField(op_mass, "q", ceed_data->spanstats.elem_restr_parent_stats, ceed_data->spanstats.basis_stats, CEED_VECTOR_ACTIVE);
200   CeedOperatorSetField(op_mass, "qdata", ceed_data->spanstats.elem_restr_parent_qd, CEED_BASIS_COLLOCATED, ceed_data->spanstats.q_data);
201   CeedOperatorSetField(op_mass, "v", ceed_data->spanstats.elem_restr_parent_stats, ceed_data->spanstats.basis_stats, CEED_VECTOR_ACTIVE);
202 
203   // Setup KSP for L^2 projection
204   {
205     MatopApplyContext M_ctx;
206     PetscInt          l_size, g_size;
207     Mat               mat_mass;
208     VecType           vec_type;
209     KSP               ksp;
210     Vec               ones, M_inv;
211     CeedVector        x_ceed, y_ceed;
212 
213     PetscCall(DMCreateGlobalVector(user->spanstats.dm, &M_inv));
214     PetscCall(VecGetLocalSize(M_inv, &l_size));
215     PetscCall(VecGetSize(M_inv, &g_size));
216     PetscCall(VecGetType(M_inv, &vec_type));
217 
218     PetscCall(PetscMalloc1(1, &M_ctx));
219     PetscCall(MatCreateShell(PETSC_COMM_WORLD, l_size, l_size, g_size, g_size, M_ctx, &mat_mass));
220     PetscCall(MatShellSetOperation(mat_mass, MATOP_MULT, (void (*)(void))MatMult_Ceed));
221     PetscCall(MatShellSetOperation(mat_mass, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed));
222     PetscCall(MatShellSetVecType(mat_mass, vec_type));
223 
224     CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_stats, &x_ceed, NULL);
225     CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_stats, &y_ceed, NULL);
226 
227     PetscCall(SetupMatopApplyCtx(PETSC_COMM_WORLD, user->spanstats.dm, user->ceed, op_mass, x_ceed, y_ceed, NULL, M_ctx));
228     user->spanstats.M_ctx = M_ctx;
229 
230     // Create lumped mass matrix inverse
231     PetscCall(DMGetGlobalVector(user->spanstats.dm, &ones));
232     PetscCall(VecZeroEntries(M_inv));
233     PetscCall(VecSet(ones, 1));
234     PetscCall(MatMult(mat_mass, ones, M_inv));
235     PetscCall(VecReciprocal(M_inv));
236     user->spanstats.M_inv = M_inv;
237     PetscCall(DMRestoreGlobalVector(user->spanstats.dm, &ones));
238 
239     PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
240     PetscCall(KSPSetOptionsPrefix(ksp, "spanstats_"));
241     {
242       PC pc;
243       PetscCall(KSPGetPC(ksp, &pc));
244       PetscCall(PCSetType(pc, PCJACOBI));
245       PetscCall(PCJacobiSetType(pc, PC_JACOBI_DIAGONAL));
246       PetscCall(KSPSetType(ksp, KSPCG));
247       PetscCall(KSPSetNormType(ksp, KSP_NORM_NATURAL));
248       PetscCall(KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT));
249     }
250     PetscCall(KSPSetOperators(ksp, mat_mass, mat_mass));
251     PetscCall(KSPSetFromOptions(ksp));
252     user->spanstats.ksp = ksp;
253   }
254 
255   // Cleanup
256   CeedQFunctionDestroy(&qf_mass);
257   PetscFunctionReturn(0);
258 }
259 
260 // Create CeedOperators and KSP for the statistics collection and processing
261 PetscErrorCode CreateStatisticsOperators(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem) {
262   CeedInt      num_comp_stats = user->spanstats.num_comp_stats, num_comp_x = problem->dim, num_comp_q;
263   CeedOperator op_setup_sur;
264   PetscFunctionBeginUser;
265   CeedBasisGetNumComponents(ceed_data->basis_q, &num_comp_q);
266 
267   // Create Operator for statistics collection
268   switch (user->phys->state_var) {
269     case STATEVAR_PRIMITIVE:
270       CeedQFunctionCreateInterior(ceed, 1, ChildStatsCollection_Prim, ChildStatsCollection_Prim_loc, &ceed_data->spanstats.qf_stats_collect);
271       break;
272     case STATEVAR_CONSERVATIVE:
273       CeedQFunctionCreateInterior(ceed, 1, ChildStatsCollection_Conserv, ChildStatsCollection_Conserv_loc, &ceed_data->spanstats.qf_stats_collect);
274       break;
275   }
276 
277   if (user->app_ctx->stats_test) {
278     CeedQFunctionDestroy(&ceed_data->spanstats.qf_stats_collect);
279     CeedQFunctionCreateInterior(ceed, 1, ChildStatsCollectionTest, ChildStatsCollectionTest_loc, &ceed_data->spanstats.qf_stats_collect);
280   }
281 
282   CeedQFunctionSetContext(ceed_data->spanstats.qf_stats_collect, problem->apply_vol_ifunction.qfunction_context);
283   CeedQFunctionAddInput(ceed_data->spanstats.qf_stats_collect, "q", num_comp_q, CEED_EVAL_INTERP);
284   CeedQFunctionAddInput(ceed_data->spanstats.qf_stats_collect, "q_data", problem->q_data_size_vol, CEED_EVAL_NONE);
285   CeedQFunctionAddInput(ceed_data->spanstats.qf_stats_collect, "x", num_comp_x, CEED_EVAL_INTERP);
286   CeedQFunctionAddOutput(ceed_data->spanstats.qf_stats_collect, "v", num_comp_stats, CEED_EVAL_NONE);
287 
288   CeedOperatorCreate(ceed, ceed_data->spanstats.qf_stats_collect, NULL, NULL, &user->spanstats.op_stats_collect);
289   CeedOperatorSetField(user->spanstats.op_stats_collect, "q", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE);
290   CeedOperatorSetField(user->spanstats.op_stats_collect, "q_data", ceed_data->elem_restr_qd_i, CEED_BASIS_COLLOCATED, ceed_data->q_data);
291   CeedOperatorSetField(user->spanstats.op_stats_collect, "x", ceed_data->elem_restr_x, ceed_data->basis_x, ceed_data->x_coord);
292   CeedOperatorSetField(user->spanstats.op_stats_collect, "v", ceed_data->spanstats.elem_restr_child_colloc, CEED_BASIS_COLLOCATED,
293                        CEED_VECTOR_ACTIVE);
294 
295   // Create Operator for L^2 projection of statistics
296   // Simply take collocated parent data (with quadrature weight already applied) and multiply by weight function.
297   // Therefore, an Identity QF is sufficient
298   CeedQFunctionCreateIdentity(ceed, num_comp_stats, CEED_EVAL_NONE, CEED_EVAL_INTERP, &ceed_data->spanstats.qf_stats_proj);
299 
300   CeedOperatorCreate(ceed, ceed_data->spanstats.qf_stats_proj, NULL, NULL, &user->spanstats.op_stats_proj);
301   CeedOperatorSetField(user->spanstats.op_stats_proj, "input", ceed_data->spanstats.elem_restr_parent_colloc, CEED_BASIS_COLLOCATED,
302                        CEED_VECTOR_ACTIVE);
303   CeedOperatorSetField(user->spanstats.op_stats_proj, "output", ceed_data->spanstats.elem_restr_parent_stats, ceed_data->spanstats.basis_stats,
304                        CEED_VECTOR_ACTIVE);
305 
306   // Get q_data for lumped mass matrix formation
307   CeedOperatorCreate(ceed, ceed_data->qf_setup_sur, NULL, NULL, &op_setup_sur);
308   CeedOperatorSetField(op_setup_sur, "dx", ceed_data->spanstats.elem_restr_parent_x, ceed_data->spanstats.basis_x, CEED_VECTOR_ACTIVE);
309   CeedOperatorSetField(op_setup_sur, "weight", CEED_ELEMRESTRICTION_NONE, ceed_data->spanstats.basis_x, CEED_VECTOR_NONE);
310   CeedOperatorSetField(op_setup_sur, "surface qdata", ceed_data->spanstats.elem_restr_parent_qd, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
311   CeedOperatorApply(op_setup_sur, ceed_data->spanstats.x_coord, ceed_data->spanstats.q_data, CEED_REQUEST_IMMEDIATE);
312 
313   CeedOperatorDestroy(&op_setup_sur);
314   PetscFunctionReturn(0);
315 }
316 
317 PetscErrorCode SetupErrorTesting(Ceed ceed, User user, CeedData ceed_data) {
318   CeedInt       num_comp_stats = user->spanstats.num_comp_stats, num_comp_x;
319   CeedQFunction qf_error;
320   CeedOperator  op_error;
321   CeedInt       q_data_size;
322   CeedVector    x_ceed, y_ceed;
323   PetscFunctionBeginUser;
324 
325   CeedElemRestrictionGetNumComponents(ceed_data->spanstats.elem_restr_parent_qd, &q_data_size);
326   CeedBasisGetNumComponents(ceed_data->spanstats.basis_x, &num_comp_x);
327 
328   CeedQFunctionCreateInterior(ceed, 1, ChildStatsCollectionTest_Error, ChildStatsCollectionTest_Error_loc, &qf_error);
329   CeedQFunctionAddInput(qf_error, "q", num_comp_stats, CEED_EVAL_INTERP);
330   CeedQFunctionAddInput(qf_error, "qdata", q_data_size, CEED_EVAL_NONE);
331   CeedQFunctionAddInput(qf_error, "x", num_comp_x, CEED_EVAL_INTERP);
332   CeedQFunctionAddOutput(qf_error, "v", num_comp_stats, CEED_EVAL_INTERP);
333 
334   CeedOperatorCreate(ceed, qf_error, NULL, NULL, &op_error);
335   CeedOperatorSetField(op_error, "q", ceed_data->spanstats.elem_restr_parent_stats, ceed_data->spanstats.basis_stats, CEED_VECTOR_ACTIVE);
336   CeedOperatorSetField(op_error, "qdata", ceed_data->spanstats.elem_restr_parent_qd, CEED_BASIS_COLLOCATED, ceed_data->spanstats.q_data);
337   CeedOperatorSetField(op_error, "x", ceed_data->spanstats.elem_restr_parent_x, ceed_data->spanstats.basis_x, ceed_data->spanstats.x_coord);
338   CeedOperatorSetField(op_error, "v", ceed_data->spanstats.elem_restr_parent_stats, ceed_data->spanstats.basis_stats, CEED_VECTOR_ACTIVE);
339 
340   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_stats, &x_ceed, NULL);
341   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_stats, &y_ceed, NULL);
342 
343   PetscCall(PetscCalloc1(1, &user->spanstats.test_error_ctx));
344   PetscCall(SetupMatopApplyCtx(PETSC_COMM_WORLD, user->spanstats.dm, user->ceed, op_error, x_ceed, y_ceed, NULL, user->spanstats.test_error_ctx));
345 
346   PetscFunctionReturn(0);
347 }
348 
349 // Setup for statistics collection
350 PetscErrorCode SetupStatsCollection(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem) {
351   DM                 dm   = user->spanstats.dm;
352   MPI_Comm           comm = PetscObjectComm((PetscObject)dm);
353   CeedInt            dim, P, Q, num_comp_x;
354   Vec                X_loc;
355   PetscMemType       X_loc_memtype;
356   const PetscScalar *X_loc_array;
357   PetscFunctionBeginUser;
358 
359   PetscCall(DMGetDimension(dm, &dim));
360   CeedBasisGetNumQuadraturePoints1D(ceed_data->basis_q, &Q);
361   CeedBasisGetNumNodes1D(ceed_data->basis_q, &P);
362 
363   PetscCall(GetRestrictionForDomain(ceed, dm, 0, 0, 0, Q, problem->q_data_size_sur, &ceed_data->spanstats.elem_restr_parent_stats,
364                                     &ceed_data->spanstats.elem_restr_parent_x, &ceed_data->spanstats.elem_restr_parent_qd));
365   CeedElemRestrictionGetNumComponents(ceed_data->spanstats.elem_restr_parent_x, &num_comp_x);
366   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_x, &ceed_data->spanstats.x_coord, NULL);
367   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_stats, &user->spanstats.rhs_ceed, NULL);
368   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_parent_qd, &ceed_data->spanstats.q_data, NULL);
369 
370   CeedBasisCreateTensorH1Lagrange(ceed, dim, num_comp_x, 2, Q, CEED_GAUSS, &ceed_data->spanstats.basis_x);
371   CeedBasisCreateTensorH1Lagrange(ceed, dim, user->spanstats.num_comp_stats, P, Q, CEED_GAUSS, &ceed_data->spanstats.basis_stats);
372 
373   PetscCall(CreateElemRestrColloc(ceed, user->spanstats.num_comp_stats, ceed_data->spanstats.basis_stats,
374                                   ceed_data->spanstats.elem_restr_parent_stats, &ceed_data->spanstats.elem_restr_parent_colloc,
375                                   &user->spanstats.parent_stats, NULL));
376   PetscCall(CreateElemRestrColloc(ceed, user->spanstats.num_comp_stats, ceed_data->basis_q, ceed_data->elem_restr_q,
377                                   &ceed_data->spanstats.elem_restr_child_colloc, &user->spanstats.child_stats, NULL));
378   CeedElemRestrictionCreateVector(ceed_data->spanstats.elem_restr_child_colloc, &user->spanstats.child_inst_stats, NULL);
379   CeedVectorSetValue(user->spanstats.child_stats, 0);
380 
381   // -- Copy DM coordinates into CeedVector
382   {
383     DM cdm;
384     PetscCall(DMGetCellCoordinateDM(dm, &cdm));
385     if (cdm) {
386       PetscCall(DMGetCellCoordinatesLocal(dm, &X_loc));
387     } else {
388       PetscCall(DMGetCoordinatesLocal(dm, &X_loc));
389     }
390   }
391   PetscCall(VecScale(X_loc, problem->dm_scale));
392   PetscCall(VecGetArrayReadAndMemType(X_loc, &X_loc_array, &X_loc_memtype));
393   CeedVectorSetArray(ceed_data->spanstats.x_coord, MemTypeP2C(X_loc_memtype), CEED_COPY_VALUES, (PetscScalar *)X_loc_array);
394   PetscCall(VecRestoreArrayRead(X_loc, &X_loc_array));
395 
396   // Create SF for communicating child data back their respective parents
397   PetscCall(PetscSFCreate(comm, &user->spanstats.sf));
398   PetscCall(CreateStatsSF(ceed, ceed_data, user->dm, user->spanstats.dm, user->spanstats.sf));
399 
400   // Create CeedOperators for statistics collection
401   PetscCall(CreateStatisticsOperators(ceed, user, ceed_data, problem));
402 
403   // Setup KSP and Mat for L^2 projection of statistics
404   PetscCall(SetupL2ProjectionStats(ceed, user, ceed_data));
405 
406   if (user->app_ctx->stats_test) {
407     PetscCall(SetupErrorTesting(ceed, user, ceed_data));
408   }
409 
410   // Setup stats viewer with prefix
411   {
412     PetscViewerType viewer_type;
413     PetscCall(PetscViewerGetType(user->app_ctx->stats_viewer, &viewer_type));
414     PetscCall(PetscOptionsSetValue(NULL, "-stats_viewer_type", viewer_type));
415 
416     PetscCall(PetscViewerSetOptionsPrefix(user->app_ctx->stats_viewer, "stats_"));
417     PetscCall(PetscViewerSetFromOptions(user->app_ctx->stats_viewer));
418   }
419   PetscFunctionReturn(0);
420 }
421 
422 // Collect statistics based on the solution Q
423 PetscErrorCode CollectStatistics(User user, PetscScalar solution_time, Vec Q) {
424   PetscMemType       q_mem_type;
425   const PetscScalar *q_arr;
426   Vec                Q_loc;
427   PetscFunctionBeginUser;
428 
429   PetscLogStage stage_stats_collect;
430   PetscCall(PetscLogStageGetId("Stats Collect", &stage_stats_collect));
431   if (stage_stats_collect == -1) PetscCall(PetscLogStageRegister("Stats Collect", &stage_stats_collect));
432   PetscCall(PetscLogStagePush(stage_stats_collect));
433 
434   PetscCall(DMGetLocalVector(user->dm, &Q_loc));
435   PetscCall(VecZeroEntries(Q_loc));
436   PetscCall(DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc));
437 
438   PetscCall(VecGetArrayReadAndMemType(Q_loc, &q_arr, &q_mem_type));
439   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER, (PetscScalar *)q_arr);
440 
441   CeedOperatorApply(user->spanstats.op_stats_collect, user->q_ceed, user->spanstats.child_inst_stats, CEED_REQUEST_IMMEDIATE);
442 
443   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
444   PetscCall(VecRestoreArrayReadAndMemType(Q_loc, &q_arr));
445   PetscCall(DMRestoreLocalVector(user->dm, &Q_loc));
446 
447   // Record averaging using left rectangle rule
448   PetscScalar delta_t           = solution_time - user->spanstats.prev_time;
449   PetscScalar prev_timeinterval = user->spanstats.prev_time - user->app_ctx->cont_time;
450   CeedVectorScale(user->spanstats.child_stats, prev_timeinterval / (prev_timeinterval + delta_t));
451   CeedVectorAXPY(user->spanstats.child_stats, delta_t / (prev_timeinterval + delta_t), user->spanstats.child_inst_stats);
452   user->spanstats.prev_time = solution_time;
453 
454   PetscCall(PetscLogStagePop());
455   PetscFunctionReturn(0);
456 }
457 
458 // Process the child statistics into parent statistics and project them onto stats
459 PetscErrorCode ProcessStatistics(User user, Vec *stats) {
460   Span_Stats         user_stats = user->spanstats;
461   const PetscScalar *child_stats;
462   PetscScalar       *parent_stats;
463   MPI_Datatype       unit;
464   Vec                rhs_loc, rhs;
465   PetscMemType       rhs_mem_type;
466   CeedScalar        *rhs_arr;
467   CeedMemType        ceed_mem_type;
468   PetscFunctionBeginUser;
469 
470   PetscLogStage stage_stats_process;
471   PetscCall(PetscLogStageGetId("Stats Process", &stage_stats_process));
472   if (stage_stats_process == -1) PetscCall(PetscLogStageRegister("Stats Process", &stage_stats_process));
473   PetscCall(PetscLogStagePush(stage_stats_process));
474 
475   CeedGetPreferredMemType(user->ceed, &ceed_mem_type);
476   CeedVectorSetValue(user_stats.parent_stats, 0);
477 
478   CeedVectorGetArrayRead(user_stats.child_stats, ceed_mem_type, &child_stats);
479   CeedVectorGetArray(user_stats.parent_stats, ceed_mem_type, &parent_stats);
480 
481   if (user_stats.num_comp_stats == 1) unit = MPIU_REAL;
482   else {
483     PetscCallMPI(MPI_Type_contiguous(user_stats.num_comp_stats, MPIU_REAL, &unit));
484     PetscCallMPI(MPI_Type_commit(&unit));
485   }
486 
487   PetscCall(PetscSFReduceBegin(user_stats.sf, unit, child_stats, parent_stats, MPI_SUM));
488   PetscCall(PetscSFReduceEnd(user_stats.sf, unit, child_stats, parent_stats, MPI_SUM));
489 
490   CeedVectorRestoreArrayRead(user_stats.child_stats, &child_stats);
491   CeedVectorRestoreArray(user_stats.parent_stats, &parent_stats);
492   PetscCallMPI(MPI_Type_free(&unit));
493 
494   CeedVectorScale(user_stats.parent_stats, 1 / user_stats.span_width);
495 
496   // L^2 projection with the parent_data
497   PetscCall(DMGetGlobalVector(user_stats.dm, &rhs));
498   PetscCall(DMGetLocalVector(user_stats.dm, &rhs_loc));
499   PetscCall(VecZeroEntries(rhs));
500   PetscCall(VecZeroEntries(rhs_loc));
501   PetscCall(VecGetArrayWriteAndMemType(rhs_loc, &rhs_arr, &rhs_mem_type));
502   CeedVectorSetArray(user_stats.rhs_ceed, MemTypeP2C(rhs_mem_type), CEED_USE_POINTER, (PetscScalar *)rhs_arr);
503 
504   CeedOperatorApply(user_stats.op_stats_proj, user_stats.parent_stats, user_stats.rhs_ceed, CEED_REQUEST_IMMEDIATE);
505 
506   CeedVectorTakeArray(user_stats.rhs_ceed, MemTypeP2C(rhs_mem_type), &rhs_arr);
507   PetscCall(VecRestoreArrayAndMemType(rhs_loc, &rhs_arr));
508   PetscCall(DMLocalToGlobal(user_stats.dm, rhs_loc, ADD_VALUES, rhs));
509 
510   PetscCall(VecDuplicate(rhs, stats));
511   PetscCall(VecPointwiseMult(*stats, rhs, user_stats.M_inv));
512 
513   PetscCall(KSPSolve(user_stats.ksp, rhs, *stats));
514 
515   PetscCall(PetscLogStagePop());
516   PetscFunctionReturn(0);
517 }
518 
519 // TSMonitor for the statistics collection and processing
520 PetscErrorCode TSMonitor_Statistics(TS ts, PetscInt steps, PetscReal solution_time, Vec Q, void *ctx) {
521   User user = (User)ctx;
522   Vec  stats;
523   PetscFunctionBeginUser;
524 
525   // Do not collect or process on the first step of the run (ie. on the initial condition)
526   if (steps == user->app_ctx->cont_steps && !user->spanstats.monitor_final_call) PetscFunctionReturn(0);
527 
528   if (steps % user->app_ctx->stats_collect_interval == 0 || user->spanstats.monitor_final_call) {
529     PetscCall(CollectStatistics(user, solution_time, Q));
530   }
531 
532   if ((steps % user->app_ctx->stats_viewer_interval == 0 && user->app_ctx->stats_viewer_interval != -1) || user->spanstats.monitor_final_call) {
533     PetscCall(DMGetGlobalVector(user->spanstats.dm, &stats));
534     PetscCall(ProcessStatistics(user, &stats));
535     PetscCall(DMSetOutputSequenceNumber(user->spanstats.dm, steps, solution_time));
536     PetscCall(PetscViewerPushFormat(user->app_ctx->stats_viewer, user->app_ctx->stats_viewer_format));
537     PetscCall(VecView(stats, user->app_ctx->stats_viewer));
538     PetscCall(PetscViewerPopFormat(user->app_ctx->stats_viewer));
539     if (user->app_ctx->stats_test && user->spanstats.monitor_final_call) {
540       Vec error;
541       PetscCall(VecDuplicate(stats, &error));
542       PetscCall(ApplyLocal_Ceed(stats, error, user->spanstats.test_error_ctx));
543       PetscScalar error_sq = 0;
544       PetscCall(VecSum(error, &error_sq));
545       PetscScalar l2_error = sqrt(error_sq);
546       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "l2 error: %.5e\n", l2_error));
547     }
548     PetscCall(DMRestoreGlobalVector(user->spanstats.dm, &stats));
549   }
550   PetscFunctionReturn(0);
551 }
552 
553 PetscErrorCode CleanupStats(User user, CeedData ceed_data) {
554   PetscFunctionBeginUser;
555 
556   // -- CeedVectors
557   CeedVectorDestroy(&user->spanstats.child_stats);
558   CeedVectorDestroy(&user->spanstats.child_inst_stats);
559   CeedVectorDestroy(&user->spanstats.parent_stats);
560   CeedVectorDestroy(&user->spanstats.rhs_ceed);
561   CeedVectorDestroy(&user->spanstats.x_ceed);
562   CeedVectorDestroy(&user->spanstats.y_ceed);
563   CeedVectorDestroy(&ceed_data->spanstats.x_coord);
564   CeedVectorDestroy(&ceed_data->spanstats.q_data);
565   CeedVectorDestroy(&user->spanstats.M_ctx->x_ceed);
566   CeedVectorDestroy(&user->spanstats.M_ctx->y_ceed);
567   if (user->app_ctx->stats_test) {
568     CeedVectorDestroy(&user->spanstats.test_error_ctx->x_ceed);
569     CeedVectorDestroy(&user->spanstats.test_error_ctx->y_ceed);
570   }
571 
572   // -- QFunctions
573   CeedQFunctionDestroy(&ceed_data->spanstats.qf_stats_collect);
574   CeedQFunctionDestroy(&ceed_data->spanstats.qf_stats_proj);
575 
576   // -- CeedBasis
577   CeedBasisDestroy(&ceed_data->spanstats.basis_stats);
578   CeedBasisDestroy(&ceed_data->spanstats.basis_x);
579 
580   // -- CeedElemRestriction
581   CeedElemRestrictionDestroy(&ceed_data->spanstats.elem_restr_parent_stats);
582   CeedElemRestrictionDestroy(&ceed_data->spanstats.elem_restr_parent_qd);
583   CeedElemRestrictionDestroy(&ceed_data->spanstats.elem_restr_parent_x);
584   CeedElemRestrictionDestroy(&ceed_data->spanstats.elem_restr_parent_colloc);
585   CeedElemRestrictionDestroy(&ceed_data->spanstats.elem_restr_child_colloc);
586 
587   // -- CeedOperators
588   CeedOperatorDestroy(&user->spanstats.op_stats_collect);
589   CeedOperatorDestroy(&user->spanstats.op_stats_proj);
590   CeedOperatorDestroy(&user->spanstats.M_ctx->op);
591   if (user->app_ctx->stats_test) CeedOperatorDestroy(&user->spanstats.test_error_ctx->op);
592 
593   // -- Vec
594   PetscCall(VecDestroy(&user->spanstats.M_inv));
595 
596   // -- KSP
597   PetscCall(KSPDestroy(&user->spanstats.ksp));
598 
599   // -- SF
600   PetscCall(PetscSFDestroy(&user->spanstats.sf));
601 
602   // -- DM
603   PetscCall(DMDestroy(&user->spanstats.dm));
604 
605   PetscFunctionReturn(0);
606 }
607