xref: /honee/problems/sgs_dd_model.c (revision 0ded6df58d56e1a7f5d8f32a7927590fb59dd6bf)
1 // Copyright (c) 2017-2024, 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 "../qfunctions/sgs_dd_model.h"
9 
10 #include <petscdmplex.h>
11 
12 #include <navierstokes.h>
13 #include <sgs_model_torch.h>
14 
15 typedef struct {
16   CeedElemRestriction      elem_restr_grid_aniso, elem_restr_sgs;
17   CeedVector               grid_aniso_ceed;
18   CeedQFunctionContext     sgsdd_qfctx, ifunction_qfctx;
19   SGSModelDDImplementation sgs_dd_model_implementation;
20 } *SgsDDSetupData;
21 
22 PetscErrorCode SgsDDSetupDataDestroy(SgsDDSetupData sgs_dd_setup_data) {
23   Ceed ceed;
24 
25   PetscFunctionBeginUser;
26   PetscCall(CeedElemRestrictionGetCeed(sgs_dd_setup_data->elem_restr_sgs, &ceed));
27 
28   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_grid_aniso));
29   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_sgs));
30   PetscCallCeed(ceed, CeedVectorDestroy(&sgs_dd_setup_data->grid_aniso_ceed));
31   PetscCallCeed(ceed, CeedQFunctionContextDestroy(&sgs_dd_setup_data->sgsdd_qfctx));
32   PetscCallCeed(ceed, CeedQFunctionContextDestroy(&sgs_dd_setup_data->ifunction_qfctx));
33   PetscCall(PetscFree(sgs_dd_setup_data));
34   PetscFunctionReturn(PETSC_SUCCESS);
35 }
36 
37 // @brief Create DM for storing subgrid stress at nodes
38 static PetscErrorCode SgsDDCreateDM(DM dm_source, DM *dm_sgs, PetscInt degree, PetscInt q_extra, PetscInt *num_components) {
39   PetscSection section;
40 
41   PetscFunctionBeginUser;
42   *num_components = 6;
43 
44   PetscCall(DMClone(dm_source, dm_sgs));
45   PetscCall(PetscObjectSetName((PetscObject)*dm_sgs, "Subgrid Stress Projection"));
46 
47   PetscCall(DMSetupByOrder_FEM(PETSC_TRUE, PETSC_TRUE, degree, 1, q_extra, 1, num_components, *dm_sgs));
48 
49   PetscCall(DMGetLocalSection(*dm_sgs, &section));
50   PetscCall(PetscSectionSetFieldName(section, 0, ""));
51   PetscCall(PetscSectionSetComponentName(section, 0, 0, "KMSubgridStressXX"));
52   PetscCall(PetscSectionSetComponentName(section, 0, 1, "KMSubgridStressYY"));
53   PetscCall(PetscSectionSetComponentName(section, 0, 2, "KMSubgridStressZZ"));
54   PetscCall(PetscSectionSetComponentName(section, 0, 3, "KMSubgridStressYZ"));
55   PetscCall(PetscSectionSetComponentName(section, 0, 4, "KMSubgridStressXZ"));
56   PetscCall(PetscSectionSetComponentName(section, 0, 5, "KMSubgridStressXY"));
57   PetscFunctionReturn(PETSC_SUCCESS);
58 };
59 
60 // @brief Evaluate data-driven SGS using fused method
61 static PetscErrorCode SgsDDNodalStressEval_Fused(User user, Vec Q_loc, Vec VelocityGradient, Vec SGSNodal_loc) {
62   SgsDDData    sgs_dd_data = user->sgs_dd_data;
63   PetscMemType q_mem_type;
64 
65   PetscFunctionBeginUser;
66   PetscCall(VecPetscToCeed(Q_loc, &q_mem_type, user->q_ceed));  // q_ceed is an implicit input
67 
68   PetscCall(ApplyCeedOperatorGlobalToLocal(VelocityGradient, SGSNodal_loc, sgs_dd_data->op_nodal_evaluation_ctx));
69 
70   PetscCall(VecCeedToPetsc(user->q_ceed, q_mem_type, Q_loc));
71   PetscFunctionReturn(PETSC_SUCCESS);
72 }
73 
74 // @brief Create CeedOperator to calculate data-drive SGS at nodes using fused operator
75 static PetscErrorCode SgsDDSetupNodalEvaluation_Fused(Ceed ceed, User user, CeedData ceed_data, SgsDDSetupData sgs_dd_setup_data) {
76   SgsDDData           sgs_dd_data = user->sgs_dd_data;
77   CeedQFunction       qf_sgs_dd_nodal;
78   CeedOperator        op_sgs_dd_nodal;
79   CeedInt             num_comp_q, num_comp_grad_velo, num_comp_x, num_comp_grid_aniso;
80   PetscInt            dim;
81   CeedVector          inv_multiplicity;
82   CeedElemRestriction elem_restr_inv_multiplicity, elem_restr_grad_velo, elem_restr_sgs;
83   DMLabel             domain_label = NULL;
84   PetscInt            label_value = 0, height = 0, dm_field = 0;
85 
86   PetscFunctionBeginUser;
87   PetscCall(DMGetDimension(user->dm, &dim));
88   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_x, &num_comp_x));
89   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q));
90   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(sgs_dd_setup_data->elem_restr_grid_aniso, &num_comp_grid_aniso));
91 
92   {  // Get velocity gradient information
93     CeedOperatorField op_field;
94     PetscCallCeed(ceed, CeedOperatorGetFieldByName(user->grad_velo_proj->l2_rhs_ctx->op, "velocity gradient", &op_field));
95     PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_grad_velo));
96     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_grad_velo, &num_comp_grad_velo));
97   }
98   PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_data->dm_sgs, domain_label, label_value, height, dm_field, &elem_restr_sgs));
99   PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_sgs, &sgs_dd_data->sgs_nodal_ceed, NULL));
100 
101   PetscCall(GetInverseMultiplicity(ceed, sgs_dd_data->dm_sgs, domain_label, label_value, height, dm_field, PETSC_FALSE, &elem_restr_inv_multiplicity,
102                                    &inv_multiplicity));
103 
104   // -- Create operator for SGS DD model nodal evaluation
105   switch (user->phys->state_var) {
106     case STATEVAR_PRIMITIVE:
107       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Prim, ComputeSgsDDNodal_Prim_loc, &qf_sgs_dd_nodal));
108       break;
109     case STATEVAR_CONSERVATIVE:
110       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Conserv, ComputeSgsDDNodal_Conserv_loc, &qf_sgs_dd_nodal));
111       break;
112     case STATEVAR_ENTROPY:
113       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Entropy, ComputeSgsDDNodal_Entropy_loc, &qf_sgs_dd_nodal));
114       break;
115   }
116 
117   // Mesh/geometry order and solution basis order may differ, therefore must interpolate
118   CeedBasis basis_x_to_q;
119   PetscCallCeed(ceed, CeedBasisCreateProjection(ceed_data->basis_x, ceed_data->basis_q, &basis_x_to_q));
120 
121   PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_nodal, sgs_dd_setup_data->sgsdd_qfctx));
122   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_nodal, "q", num_comp_q, CEED_EVAL_NONE));
123   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_nodal, "x", num_comp_x, CEED_EVAL_INTERP));
124   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_nodal, "gradient velocity", num_comp_grad_velo, CEED_EVAL_NONE));
125   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_nodal, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE));
126   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_nodal, "inverse multiplicity", 1, CEED_EVAL_NONE));
127   PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_nodal, "km_sgs", sgs_dd_data->num_comp_sgs, CEED_EVAL_NONE));
128 
129   PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_nodal, NULL, NULL, &op_sgs_dd_nodal));
130   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "q", ceed_data->elem_restr_q, CEED_BASIS_NONE, user->q_ceed));
131   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "x", ceed_data->elem_restr_x, basis_x_to_q, ceed_data->x_coord));
132   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "gradient velocity", elem_restr_grad_velo, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
133   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "anisotropy tensor", sgs_dd_setup_data->elem_restr_grid_aniso, CEED_BASIS_NONE,
134                                            sgs_dd_setup_data->grid_aniso_ceed));
135   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_NONE, inv_multiplicity));
136   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_nodal, "km_sgs", elem_restr_sgs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
137 
138   PetscCall(OperatorApplyContextCreate(user->grad_velo_proj->dm, sgs_dd_data->dm_sgs, ceed, op_sgs_dd_nodal, NULL, sgs_dd_data->sgs_nodal_ceed, NULL,
139                                        NULL, &sgs_dd_data->op_nodal_evaluation_ctx));
140 
141   sgs_dd_setup_data->elem_restr_sgs = elem_restr_sgs;
142   sgs_dd_data->sgs_nodal_eval       = SgsDDNodalStressEval_Fused;
143 
144   PetscCallCeed(ceed, CeedVectorDestroy(&inv_multiplicity));
145   PetscCallCeed(ceed, CeedBasisDestroy(&basis_x_to_q));
146   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_inv_multiplicity));
147   PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_nodal));
148   PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_nodal));
149   PetscFunctionReturn(PETSC_SUCCESS);
150 }
151 
152 // @brief Setup data-driven model inference using libCEED native implementation
153 static PetscErrorCode SgsDDSetupNodalEvaluation_Sequential_Ceed(Ceed ceed, SgsDDData sgs_dd_data, SgsDDSetupData sgs_dd_setup_data,
154                                                                 CeedElemRestriction elem_restr_dd_inputs, CeedElemRestriction elem_restr_dd_outputs,
155                                                                 CeedElemRestriction elem_restr_inv_multiplicity, CeedVector inv_multiplicity,
156                                                                 void **ctx) {
157   CeedQFunction         qf_sgs_dd_inference;
158   CeedOperator          op_sgs_dd_inference;
159   OperatorApplyContext *op_context = (OperatorApplyContext *)ctx;
160 
161   PetscFunctionBeginUser;
162   PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Sequential_Inference, ComputeSgsDDNodal_Sequential_Inference_loc,
163                                                   &qf_sgs_dd_inference));
164 
165   PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_inference, sgs_dd_setup_data->sgsdd_qfctx));
166   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inference, "model inputs", sgs_dd_data->num_comp_inputs, CEED_EVAL_NONE));
167   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inference, "inverse multiplicity", 1, CEED_EVAL_NONE));
168   PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_inference, "model outputs", sgs_dd_data->num_comp_outputs, CEED_EVAL_NONE));
169 
170   PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_inference, NULL, NULL, &op_sgs_dd_inference));
171   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inference, "model inputs", elem_restr_dd_inputs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
172   PetscCallCeed(ceed,
173                 CeedOperatorSetField(op_sgs_dd_inference, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_NONE, inv_multiplicity));
174   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inference, "model outputs", elem_restr_dd_outputs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
175 
176   PetscCall(OperatorApplyContextCreate(sgs_dd_data->dm_dd_inputs, sgs_dd_data->dm_dd_outputs, ceed, op_sgs_dd_inference, NULL, NULL, NULL, NULL,
177                                        op_context));
178   sgs_dd_data->sgs_nodal_inference_ctx_destroy = (PetscErrorCode(*)(void *))OperatorApplyContextDestroy;
179 
180   PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_inference));
181   PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_inference));
182   PetscFunctionReturn(PETSC_SUCCESS);
183 }
184 
185 // @brief Perform data-driven model inference using libCEED native implementation
186 PetscErrorCode SgsDDNodalStressEval_Sequential_Ceed(Vec DD_Inputs_loc, Vec DD_Outputs_loc, void *ctx) {
187   OperatorApplyContext op_context = *(OperatorApplyContext *)ctx;
188 
189   PetscFunctionBeginUser;
190   PetscCall(PetscLogEventBegin(FLUIDS_SgsModelDDData, DD_Inputs_loc, DD_Outputs_loc, NULL, NULL));
191   PetscCall(PetscLogEventBegin(FLUIDS_SgsModelDDInference, DD_Inputs_loc, DD_Outputs_loc, NULL, NULL));
192   PetscCall(PetscLogGpuTimeBegin());
193   PetscCall(ApplyCeedOperatorLocalToLocal(DD_Inputs_loc, DD_Outputs_loc, op_context));
194   PetscCall(PetscLogGpuTimeEnd());
195   PetscCall(PetscLogEventEnd(FLUIDS_SgsModelDDInference, DD_Inputs_loc, DD_Outputs_loc, NULL, NULL));
196   PetscCall(PetscLogEventEnd(FLUIDS_SgsModelDDData, DD_Inputs_loc, DD_Outputs_loc, NULL, NULL));
197   PetscFunctionReturn(PETSC_SUCCESS);
198 }
199 
200 // @brief Setup data-driven model inference using libtorch
201 static PetscErrorCode SgsDDSetupNodalEvaluation_Sequential_Torch(Ceed ceed, SgsDDData sgs_dd_data, SgsDDSetupData sgs_dd_setup_data,
202                                                                  CeedElemRestriction elem_restr_dd_inputs, CeedElemRestriction elem_restr_dd_outputs,
203                                                                  CeedElemRestriction elem_restr_inv_multiplicity, CeedVector inv_multiplicity,
204                                                                  void **ctx) {
205   const char     *ceed_resource;
206   char            model_path[PETSC_MAX_PATH_LEN] = "";
207   TorchDeviceType model_device_type;
208 
209   PetscFunctionBeginUser;
210   PetscCallCeed(ceed, CeedGetResource(ceed, &ceed_resource));
211   if (strstr(ceed_resource, "/gpu/cuda")) model_device_type = TORCH_DEVICE_CUDA;
212   else if (strstr(ceed_resource, "/gpu/hip")) model_device_type = TORCH_DEVICE_HIP;
213   // On-device XPU is not working reliably currently, default to CPU inference evaluation
214   // else if (strstr(ceed_resource, "/gpu/sycl")) model_device_type = TORCH_DEVICE_XPU;
215   else model_device_type = TORCH_DEVICE_CPU;
216   PetscCall(PetscOptionsGetEnum(NULL, NULL, "-sgs_model_dd_torch_model_device", TorchDeviceTypes, (PetscEnum *)&model_device_type, NULL));
217   PetscCall(PetscOptionsGetString(NULL, NULL, "-sgs_model_dd_torch_model_path", model_path, sizeof(model_path), NULL));
218 
219   PetscCall(LoadModel_Torch(model_path, model_device_type));
220 
221   PetscFunctionReturn(PETSC_SUCCESS);
222 }
223 
224 // @brief Perform data-driven model inference using libtorch
225 static PetscErrorCode SgsDDNodalStressEval_Sequential_Torch(Vec DD_Inputs_loc, Vec DD_Outputs_loc, void *ctx) {
226   static PetscBool run_through = PETSC_FALSE;
227   PetscFunctionBeginUser;
228   if (!run_through) {
229     PetscCall(VecViewFromOptions(DD_Inputs_loc, NULL, "-dd_inputs_loc_view"));
230   }
231   PetscCall(ModelInference_Torch(DD_Inputs_loc, DD_Outputs_loc));
232   if (!run_through) {
233     PetscCall(VecViewFromOptions(DD_Outputs_loc, NULL, "-dd_outputs_loc_view"));
234     run_through = PETSC_TRUE;
235   }
236   PetscFunctionReturn(PETSC_SUCCESS);
237 }
238 
239 // @brief Evaluate data-driven SGS using sequential method
240 PetscErrorCode SgsDDNodalStressEval_Sequential(User user, Vec Q_loc, Vec VelocityGradient, Vec SGSNodal_loc) {
241   SgsDDData    sgs_dd_data = user->sgs_dd_data;
242   PetscMemType q_mem_type;
243   Vec          DD_Inputs_loc, DD_Outputs_loc;
244 
245   PetscFunctionBeginUser;
246   PetscCall(DMGetLocalVector(sgs_dd_data->dm_dd_inputs, &DD_Inputs_loc));
247   PetscCall(DMGetLocalVector(sgs_dd_data->dm_dd_outputs, &DD_Outputs_loc));
248   PetscCall(VecPetscToCeed(Q_loc, &q_mem_type, user->q_ceed));  // q_ceed is an implicit input
249 
250   PetscCall(ApplyCeedOperatorGlobalToLocal(VelocityGradient, DD_Inputs_loc, sgs_dd_data->op_nodal_dd_inputs_ctx));
251   PetscCall(sgs_dd_data->sgs_nodal_inference(DD_Inputs_loc, DD_Outputs_loc, &sgs_dd_data->sgs_nodal_inference_ctx));
252   PetscCall(ApplyCeedOperatorLocalToLocal(DD_Outputs_loc, SGSNodal_loc, sgs_dd_data->op_nodal_dd_outputs_ctx));
253 
254   PetscCall(VecCeedToPetsc(user->q_ceed, q_mem_type, Q_loc));
255   PetscCall(DMRestoreLocalVector(sgs_dd_data->dm_dd_inputs, &DD_Inputs_loc));
256   PetscCall(DMRestoreLocalVector(sgs_dd_data->dm_dd_outputs, &DD_Outputs_loc));
257   PetscFunctionReturn(PETSC_SUCCESS);
258 }
259 
260 // @brief Create CeedOperator to calculate data-drive SGS at nodes using sequentially-applied operators
261 static PetscErrorCode SgsDDSetupNodalEvaluation_Sequential(Ceed ceed, User user, CeedData ceed_data, SgsDDSetupData sgs_dd_setup_data) {
262   SgsDDData           sgs_dd_data = user->sgs_dd_data;
263   CeedInt             num_comp_q, num_comp_grad_velo, num_comp_x, num_comp_grid_aniso, num_comp_eigvec = 9 + 1;
264   PetscInt            dim;
265   CeedVector          inv_multiplicity, eigvec;
266   CeedElemRestriction elem_restr_inv_multiplicity, elem_restr_grad_velo, elem_restr_sgs, elem_restr_eigvec, elem_restr_dd_inputs,
267       elem_restr_dd_outputs;
268   DMLabel  domain_label = NULL;
269   PetscInt label_value = 0, height = 0, dm_field = 0;
270 
271   PetscFunctionBeginUser;
272   {  // Create DMs for data-driven input and output values
273     PetscSection section;
274     PetscInt     degree, q_extra;
275     {  // Get degree and number of quadrature points from dm_sgs
276       PetscFE         fe;
277       PetscSpace      basis;
278       PetscQuadrature quadrature;
279       PetscInt        num_qpnts;
280       PetscCall(DMGetField(sgs_dd_data->dm_sgs, 0, NULL, (PetscObject *)&fe));
281       PetscCall(PetscFEGetBasisSpace(fe, &basis));
282       PetscCall(PetscSpaceGetDegree(basis, &degree, NULL));
283       PetscCall(PetscFEGetQuadrature(fe, &quadrature));
284       PetscCall(PetscQuadratureGetOrder(quadrature, &num_qpnts));
285       q_extra = degree - num_qpnts;
286     }
287 
288     PetscCall(DMClone(sgs_dd_data->dm_sgs, &sgs_dd_data->dm_dd_inputs));
289     PetscCall(PetscObjectSetName((PetscObject)sgs_dd_data->dm_dd_inputs, "Data-Driven Model Inputs"));
290     PetscCall(DMSetupByOrder_FEM(PETSC_TRUE, PETSC_TRUE, degree, 1, q_extra, 1, &sgs_dd_data->num_comp_inputs, sgs_dd_data->dm_dd_inputs));
291     PetscCall(DMGetLocalSection(sgs_dd_data->dm_dd_inputs, &section));
292     PetscCall(PetscSectionSetFieldName(section, 0, ""));
293     for (CeedInt i = 0; i < sgs_dd_data->num_comp_inputs; i++) {
294       char component_name[PETSC_MAX_PATH_LEN];
295 
296       PetscCall(PetscSNPrintf(component_name, sizeof component_name, "DataDrivenInput%" CeedInt_FMT, i + 1));
297       PetscCall(PetscSectionSetComponentName(section, 0, i, component_name));
298     }
299 
300     PetscCall(DMClone(sgs_dd_data->dm_sgs, &sgs_dd_data->dm_dd_outputs));
301     PetscCall(PetscObjectSetName((PetscObject)sgs_dd_data->dm_dd_outputs, "Data-Driven Model Outputs"));
302     PetscCall(DMSetupByOrder_FEM(PETSC_TRUE, PETSC_TRUE, degree, 1, q_extra, 1, &sgs_dd_data->num_comp_outputs, sgs_dd_data->dm_dd_outputs));
303     PetscCall(DMGetLocalSection(sgs_dd_data->dm_dd_outputs, &section));
304     PetscCall(PetscSectionSetFieldName(section, 0, ""));
305     for (CeedInt i = 0; i < sgs_dd_data->num_comp_outputs; i++) {
306       char component_name[PETSC_MAX_PATH_LEN];
307 
308       PetscCall(PetscSNPrintf(component_name, sizeof component_name, "DataDrivenOutput%" CeedInt_FMT, i + 1));
309       PetscCall(PetscSectionSetComponentName(section, 0, i, component_name));
310     }
311   }
312 
313   PetscCall(DMGetDimension(user->dm, &dim));
314   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_x, &num_comp_x));
315   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q));
316   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(sgs_dd_setup_data->elem_restr_grid_aniso, &num_comp_grid_aniso));
317 
318   {  // Get velocity gradient information
319     CeedOperatorField op_field;
320     PetscCallCeed(ceed, CeedOperatorGetFieldByName(user->grad_velo_proj->l2_rhs_ctx->op, "velocity gradient", &op_field));
321     PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_grad_velo));
322     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_grad_velo, &num_comp_grad_velo));
323     PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_grad_velo, &sgs_dd_data->grad_velo_ceed, NULL));
324   }
325 
326   PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_data->dm_sgs, domain_label, label_value, height, dm_field, &elem_restr_sgs));
327   PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_sgs, &sgs_dd_data->sgs_nodal_ceed, NULL));
328   PetscCall(
329       DMPlexCeedElemRestrictionCollocatedCreate(ceed, sgs_dd_data->dm_sgs, domain_label, label_value, height, num_comp_eigvec, &elem_restr_eigvec));
330   PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_eigvec, &eigvec, NULL));
331 
332   PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_data->dm_dd_inputs, domain_label, label_value, height, dm_field, &elem_restr_dd_inputs));
333   PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_data->dm_dd_outputs, domain_label, label_value, height, dm_field, &elem_restr_dd_outputs));
334 
335   PetscCall(GetInverseMultiplicity(ceed, sgs_dd_data->dm_sgs, domain_label, label_value, height, dm_field, PETSC_FALSE, &elem_restr_inv_multiplicity,
336                                    &inv_multiplicity));
337 
338   {  // Create operator for data-driven input evaluation
339     CeedQFunction qf_sgs_dd_inputs;
340     CeedOperator  op_sgs_dd_inputs;
341 
342     switch (user->phys->state_var) {
343       case STATEVAR_PRIMITIVE:
344         PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Sequential_Inputs_Prim,
345                                                         ComputeSgsDDNodal_Sequential_Inputs_Prim_loc, &qf_sgs_dd_inputs));
346         break;
347       case STATEVAR_CONSERVATIVE:
348         PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Sequential_Inputs_Conserv,
349                                                         ComputeSgsDDNodal_Sequential_Inputs_Conserv_loc, &qf_sgs_dd_inputs));
350         break;
351       case STATEVAR_ENTROPY:
352         PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Sequential_Inputs_Entropy,
353                                                         ComputeSgsDDNodal_Sequential_Inputs_Entropy_loc, &qf_sgs_dd_inputs));
354         break;
355     }
356 
357     PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_inputs, sgs_dd_setup_data->sgsdd_qfctx));
358     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inputs, "q", num_comp_q, CEED_EVAL_NONE));
359     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inputs, "gradient velocity", num_comp_grad_velo, CEED_EVAL_NONE));
360     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inputs, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE));
361     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_inputs, "inverse multiplicity", 1, CEED_EVAL_NONE));
362     PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_inputs, "eigenvectors", num_comp_eigvec, CEED_EVAL_NONE));
363     PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_inputs, "model inputs", sgs_dd_data->num_comp_inputs, CEED_EVAL_NONE));
364 
365     PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_inputs, NULL, NULL, &op_sgs_dd_inputs));
366     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inputs, "q", ceed_data->elem_restr_q, CEED_BASIS_NONE, user->q_ceed));
367     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inputs, "gradient velocity", elem_restr_grad_velo, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
368     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inputs, "anisotropy tensor", sgs_dd_setup_data->elem_restr_grid_aniso, CEED_BASIS_NONE,
369                                              sgs_dd_setup_data->grid_aniso_ceed));
370     PetscCallCeed(ceed,
371                   CeedOperatorSetField(op_sgs_dd_inputs, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_NONE, inv_multiplicity));
372     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inputs, "eigenvectors", elem_restr_eigvec, CEED_BASIS_NONE, eigvec));
373     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_inputs, "model inputs", elem_restr_dd_inputs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
374 
375     PetscCall(OperatorApplyContextCreate(user->grad_velo_proj->dm, sgs_dd_data->dm_dd_inputs, ceed, op_sgs_dd_inputs, NULL, NULL, NULL, NULL,
376                                          &sgs_dd_data->op_nodal_dd_inputs_ctx));
377     PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_inputs));
378     PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_inputs));
379   }
380 
381   {  // Create operator for data-driven output handling
382     CeedQFunction qf_sgs_dd_outputs;
383     CeedOperator  op_sgs_dd_outputs;
384 
385     PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSgsDDNodal_Sequential_Outputs, ComputeSgsDDNodal_Sequential_Outputs_loc,
386                                                     &qf_sgs_dd_outputs));
387     PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_outputs, sgs_dd_setup_data->sgsdd_qfctx));
388     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_outputs, "model outputs", sgs_dd_data->num_comp_outputs, CEED_EVAL_NONE));
389     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_outputs, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE));
390     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_outputs, "inverse multiplicity", 1, CEED_EVAL_NONE));
391     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_outputs, "eigenvectors", num_comp_eigvec, CEED_EVAL_NONE));
392     PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_outputs, "km_sgs", sgs_dd_data->num_comp_sgs, CEED_EVAL_NONE));
393 
394     PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_outputs, NULL, NULL, &op_sgs_dd_outputs));
395     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_outputs, "model outputs", elem_restr_dd_outputs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
396     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_outputs, "anisotropy tensor", sgs_dd_setup_data->elem_restr_grid_aniso, CEED_BASIS_NONE,
397                                              sgs_dd_setup_data->grid_aniso_ceed));
398     PetscCallCeed(ceed,
399                   CeedOperatorSetField(op_sgs_dd_outputs, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_NONE, inv_multiplicity));
400     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_outputs, "eigenvectors", elem_restr_eigvec, CEED_BASIS_NONE, eigvec));
401     PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_outputs, "km_sgs", elem_restr_sgs, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE));
402 
403     PetscCall(OperatorApplyContextCreate(sgs_dd_data->dm_dd_outputs, sgs_dd_data->dm_sgs, ceed, op_sgs_dd_outputs, NULL, sgs_dd_data->sgs_nodal_ceed,
404                                          NULL, NULL, &sgs_dd_data->op_nodal_dd_outputs_ctx));
405     PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_outputs));
406     PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_outputs));
407   }
408 
409   sgs_dd_data->sgs_nodal_eval = SgsDDNodalStressEval_Sequential;
410 
411   if (sgs_dd_setup_data->sgs_dd_model_implementation == SGS_MODEL_DD_SEQENTIAL_CEED) {
412     sgs_dd_data->sgs_nodal_inference = SgsDDNodalStressEval_Sequential_Ceed;
413     PetscCall(SgsDDSetupNodalEvaluation_Sequential_Ceed(ceed, sgs_dd_data, sgs_dd_setup_data, elem_restr_dd_inputs, elem_restr_dd_outputs,
414                                                         elem_restr_inv_multiplicity, inv_multiplicity, &sgs_dd_data->sgs_nodal_inference_ctx));
415   } else if (sgs_dd_setup_data->sgs_dd_model_implementation == SGS_MODEL_DD_SEQENTIAL_TORCH) {
416     sgs_dd_data->sgs_nodal_inference = SgsDDNodalStressEval_Sequential_Torch;
417     PetscCall(SgsDDSetupNodalEvaluation_Sequential_Torch(ceed, sgs_dd_data, sgs_dd_setup_data, elem_restr_dd_inputs, elem_restr_dd_outputs,
418                                                          elem_restr_inv_multiplicity, inv_multiplicity, &sgs_dd_data->sgs_nodal_inference_ctx));
419   }
420 
421   sgs_dd_setup_data->elem_restr_sgs = elem_restr_sgs;
422 
423   PetscCallCeed(ceed, CeedVectorDestroy(&inv_multiplicity));
424   PetscCallCeed(ceed, CeedVectorDestroy(&eigvec));
425   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_inv_multiplicity));
426   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_eigvec));
427   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_dd_inputs));
428   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_dd_outputs));
429   PetscFunctionReturn(PETSC_SUCCESS);
430 }
431 
432 // @brief Create CeedOperator to compute SGS contribution to the residual
433 static PetscErrorCode SgsSetupNodalIFunction(Ceed ceed, User user, CeedData ceed_data, SgsDDSetupData sgs_dd_setup_data) {
434   SgsDDData     sgs_dd_data = user->sgs_dd_data;
435   CeedInt       num_comp_q, num_comp_qd, num_comp_x;
436   PetscInt      dim;
437   CeedQFunction qf_sgs_apply;
438   CeedOperator  op_sgs_apply;
439   CeedBasis     basis_sgs;
440 
441   PetscFunctionBeginUser;
442   PetscCall(DMGetDimension(user->dm, &dim));
443   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q));
444   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &num_comp_qd));
445   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_x, &num_comp_x));
446 
447   PetscCall(CreateBasisFromPlex(ceed, sgs_dd_data->dm_sgs, 0, 0, 0, 0, &basis_sgs));
448 
449   switch (user->phys->state_var) {
450     case STATEVAR_PRIMITIVE:
451       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, IFunction_NodalSgs_Prim, IFunction_NodalSgs_Prim_loc, &qf_sgs_apply));
452       break;
453     case STATEVAR_CONSERVATIVE:
454       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, IFunction_NodalSgs_Conserv, IFunction_NodalSgs_Conserv_loc, &qf_sgs_apply));
455       break;
456     case STATEVAR_ENTROPY:
457       PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, IFunction_NodalSgs_Entropy, IFunction_NodalSgs_Entropy_loc, &qf_sgs_apply));
458       break;
459   }
460 
461   PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_apply, sgs_dd_setup_data->ifunction_qfctx));
462   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_apply, "q", num_comp_q, CEED_EVAL_INTERP));
463   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_apply, "qdata", num_comp_qd, CEED_EVAL_NONE));
464   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_apply, "km_sgs", sgs_dd_data->num_comp_sgs, CEED_EVAL_INTERP));
465   PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_apply, "Grad_v", num_comp_q * dim, CEED_EVAL_GRAD));
466 
467   PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_apply, NULL, NULL, &op_sgs_apply));
468   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_apply, "q", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE));
469   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_apply, "qdata", ceed_data->elem_restr_qd_i, CEED_BASIS_NONE, ceed_data->q_data));
470   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_apply, "km_sgs", sgs_dd_setup_data->elem_restr_sgs, basis_sgs, sgs_dd_data->sgs_nodal_ceed));
471   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_apply, "Grad_v", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE));
472 
473   PetscCall(
474       OperatorApplyContextCreate(user->dm, user->dm, ceed, op_sgs_apply, user->q_ceed, user->g_ceed, NULL, NULL, &sgs_dd_data->op_sgs_apply_ctx));
475 
476   PetscCallCeed(ceed, CeedBasisDestroy(&basis_sgs));
477   PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_apply));
478   PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_apply));
479   PetscFunctionReturn(PETSC_SUCCESS);
480 }
481 
482 // @brief Calculate and add data-driven SGS residual to the global residual
483 PetscErrorCode SgsDDApplyIFunction(User user, const Vec Q_loc, Vec G_loc) {
484   SgsDDData    sgs_dd_data = user->sgs_dd_data;
485   Vec          VelocityGradient, SGSNodal_loc;
486   PetscMemType sgs_nodal_mem_type;
487 
488   PetscFunctionBeginUser;
489   PetscCall(PetscLogEventBegin(FLUIDS_SgsModel, Q_loc, G_loc, NULL, NULL));
490   PetscCall(DMGetGlobalVector(user->grad_velo_proj->dm, &VelocityGradient));
491   PetscCall(VelocityGradientProjectionApply(user->grad_velo_proj, Q_loc, VelocityGradient));
492 
493   // -- Compute Nodal SGS tensor
494   PetscCall(DMGetLocalVector(sgs_dd_data->dm_sgs, &SGSNodal_loc));
495   PetscCall(sgs_dd_data->sgs_nodal_eval(user, Q_loc, VelocityGradient, SGSNodal_loc));
496 
497   // -- Compute contribution of the SGS stress
498   PetscCall(VecPetscToCeed(SGSNodal_loc, &sgs_nodal_mem_type, sgs_dd_data->sgs_nodal_ceed));  // sgs_nodal_ceed is an implicit input
499   PetscCall(ApplyAddCeedOperatorLocalToLocal(Q_loc, G_loc, sgs_dd_data->op_sgs_apply_ctx));
500 
501   // -- Return local SGS vector
502   PetscCall(VecCeedToPetsc(sgs_dd_data->sgs_nodal_ceed, sgs_nodal_mem_type, SGSNodal_loc));
503   PetscCall(DMRestoreLocalVector(sgs_dd_data->dm_sgs, &SGSNodal_loc));
504   PetscCall(DMRestoreGlobalVector(user->grad_velo_proj->dm, &VelocityGradient));
505   PetscCall(PetscLogEventEnd(FLUIDS_SgsModel, Q_loc, G_loc, NULL, NULL));
506   PetscFunctionReturn(PETSC_SUCCESS);
507 }
508 
509 // @brief B = A^T, A is NxM, B is MxN
510 static PetscErrorCode TransposeMatrix(const PetscScalar *A, PetscScalar *B, const PetscInt N, const PetscInt M) {
511   PetscFunctionBeginUser;
512   for (PetscInt i = 0; i < N; i++) {
513     for (PetscInt j = 0; j < M; j++) {
514       B[j * N + i] = A[i * M + j];
515     }
516   }
517   PetscFunctionReturn(PETSC_SUCCESS);
518 }
519 
520 // @brief Read neural network coefficients from file and put into context struct
521 static PetscErrorCode SgsDDContextFill(MPI_Comm comm, char data_dir[PETSC_MAX_PATH_LEN], SgsDDContext *psgsdd_ctx) {
522   SgsDDContext sgsdd_ctx;
523   PetscInt     num_inputs = (*psgsdd_ctx)->num_inputs, num_outputs = (*psgsdd_ctx)->num_outputs, num_neurons = (*psgsdd_ctx)->num_neurons;
524   char         file_path[PETSC_MAX_PATH_LEN];
525   PetscScalar *temp;
526 
527   PetscFunctionBeginUser;
528   {
529     SgsDDContext sgsdd_temp;
530     PetscCall(PetscNew(&sgsdd_temp));
531     *sgsdd_temp                     = **psgsdd_ctx;
532     sgsdd_temp->offsets.bias1       = 0;
533     sgsdd_temp->offsets.bias2       = sgsdd_temp->offsets.bias1 + num_neurons;
534     sgsdd_temp->offsets.weight1     = sgsdd_temp->offsets.bias2 + num_neurons;
535     sgsdd_temp->offsets.weight2     = sgsdd_temp->offsets.weight1 + num_neurons * num_inputs;
536     sgsdd_temp->offsets.out_scaling = sgsdd_temp->offsets.weight2 + num_inputs * num_neurons;
537     PetscInt total_num_scalars      = sgsdd_temp->offsets.out_scaling + 2 * num_outputs;
538     sgsdd_temp->total_bytes         = sizeof(*sgsdd_ctx) + total_num_scalars * sizeof(sgsdd_ctx->data[0]);
539     PetscCall(PetscMalloc(sgsdd_temp->total_bytes, &sgsdd_ctx));
540     *sgsdd_ctx = *sgsdd_temp;
541     PetscCall(PetscFree(sgsdd_temp));
542   }
543 
544   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b1.dat"));
545   PetscCall(PhastaDatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias1]));
546   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b2.dat"));
547   PetscCall(PhastaDatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias2]));
548   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "OutScaling.dat"));
549   PetscCall(PhastaDatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling]));
550 
551   {
552     PetscCall(PetscMalloc1(num_inputs * num_neurons, &temp));
553     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w1.dat"));
554     PetscCall(PhastaDatFileReadToArrayReal(comm, file_path, temp));
555     PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight1], num_inputs, num_neurons));
556     PetscCall(PetscFree(temp));
557   }
558   {
559     PetscCall(PetscMalloc1(num_outputs * num_neurons, &temp));
560     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w2.dat"));
561     PetscCall(PhastaDatFileReadToArrayReal(comm, file_path, temp));
562     PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight2], num_neurons, num_outputs));
563     PetscCall(PetscFree(temp));
564   }
565 
566   PetscCall(PetscFree(*psgsdd_ctx));
567   *psgsdd_ctx = sgsdd_ctx;
568   PetscFunctionReturn(PETSC_SUCCESS);
569 }
570 
571 PetscErrorCode SgsDDSetup(Ceed ceed, User user, CeedData ceed_data, ProblemData problem) {
572   PetscReal                alpha = 0;
573   SgsDDContext             sgsdd_ctx;
574   MPI_Comm                 comm                           = user->comm;
575   char                     sgs_dd_dir[PETSC_MAX_PATH_LEN] = "./dd_sgs_parameters";
576   SgsDDSetupData           sgs_dd_setup_data;
577   NewtonianIdealGasContext gas;
578 
579   PetscFunctionBeginUser;
580   PetscCall(VelocityGradientProjectionSetup(ceed, user, ceed_data, problem, user->phys->state_var, ceed_data->elem_restr_q, ceed_data->basis_q,
581                                             &user->grad_velo_proj));
582 
583   PetscCall(PetscNew(&user->sgs_dd_data));
584   user->sgs_dd_data->num_comp_inputs  = 6;
585   user->sgs_dd_data->num_comp_outputs = 6;
586 
587   PetscCall(PetscNew(&sgs_dd_setup_data));
588 
589   PetscOptionsBegin(comm, NULL, "SGS Data-Driven Model Options", NULL);
590   PetscCall(PetscOptionsReal("-sgs_model_dd_leakyrelu_alpha", "Slope parameter for Leaky ReLU activation function", NULL, alpha, &alpha, NULL));
591   PetscCall(PetscOptionsString("-sgs_model_dd_parameter_dir", "Path to directory with model parameters (weights, biases, etc.)", NULL, sgs_dd_dir,
592                                sgs_dd_dir, sizeof(sgs_dd_dir), NULL));
593   PetscCall(PetscOptionsDeprecated("-sgs_model_dd_use_fused", NULL, "libCEED 0.12.0", "Use -sgs_model_dd_type instead"));
594   sgs_dd_setup_data->sgs_dd_model_implementation = SGS_MODEL_DD_FUSED;
595   PetscCall(PetscOptionsEnum("-sgs_model_dd_implementation", "Data-Driven SGS model implementation", NULL, SGSModelDDImplementations,
596                              (PetscEnum)sgs_dd_setup_data->sgs_dd_model_implementation, (PetscEnum *)&sgs_dd_setup_data->sgs_dd_model_implementation,
597                              NULL));
598   PetscOptionsEnd();
599 
600   PetscCall(PetscNew(&sgsdd_ctx));
601   sgsdd_ctx->num_layers  = 1;
602   sgsdd_ctx->num_inputs  = 6;
603   sgsdd_ctx->num_outputs = 6;
604   sgsdd_ctx->num_neurons = 20;
605   sgsdd_ctx->alpha       = alpha;
606 
607   PetscCall(SgsDDContextFill(comm, sgs_dd_dir, &sgsdd_ctx));
608 
609   // -- Create DM for storing SGS tensor at nodes
610   PetscCall(SgsDDCreateDM(user->dm, &user->sgs_dd_data->dm_sgs, user->app_ctx->degree, user->app_ctx->q_extra, &user->sgs_dd_data->num_comp_sgs));
611 
612   PetscCallCeed(ceed, CeedQFunctionContextGetDataRead(problem->apply_vol_ifunction.qfunction_context, CEED_MEM_HOST, &gas));
613   sgsdd_ctx->gas = *gas;
614   PetscCallCeed(ceed, CeedQFunctionContextRestoreDataRead(problem->apply_vol_ifunction.qfunction_context, &gas));
615   PetscCallCeed(ceed, CeedQFunctionContextCreate(user->ceed, &sgs_dd_setup_data->sgsdd_qfctx));
616   PetscCallCeed(ceed,
617                 CeedQFunctionContextSetData(sgs_dd_setup_data->sgsdd_qfctx, CEED_MEM_HOST, CEED_USE_POINTER, sgsdd_ctx->total_bytes, sgsdd_ctx));
618   PetscCallCeed(ceed, CeedQFunctionContextSetDataDestroy(sgs_dd_setup_data->sgsdd_qfctx, CEED_MEM_HOST, FreeContextPetsc));
619 
620   PetscCallCeed(ceed, CeedQFunctionContextReferenceCopy(problem->apply_vol_ifunction.qfunction_context, &sgs_dd_setup_data->ifunction_qfctx));
621 
622   // -- Compute and store anisotropy tensor
623   PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, user, ceed_data, &sgs_dd_setup_data->elem_restr_grid_aniso,
624                                                      &sgs_dd_setup_data->grid_aniso_ceed));
625 
626   // -- Create Nodal Evaluation Operator
627   switch (sgs_dd_setup_data->sgs_dd_model_implementation) {
628     case SGS_MODEL_DD_FUSED:
629       PetscCall(SgsDDSetupNodalEvaluation_Fused(ceed, user, ceed_data, sgs_dd_setup_data));
630       break;
631     case SGS_MODEL_DD_SEQENTIAL_CEED:
632     case SGS_MODEL_DD_SEQENTIAL_TORCH:
633       PetscCall(SgsDDSetupNodalEvaluation_Sequential(ceed, user, ceed_data, sgs_dd_setup_data));
634       break;
635   }
636 
637   // -- Create Operator to evalutate residual of SGS stress
638   PetscCall(SgsSetupNodalIFunction(ceed, user, ceed_data, sgs_dd_setup_data));
639 
640   PetscCall(SgsDDSetupDataDestroy(sgs_dd_setup_data));
641   PetscFunctionReturn(PETSC_SUCCESS);
642 }
643 
644 PetscErrorCode SgsDDDataDestroy(SgsDDData sgs_dd_data) {
645   PetscFunctionBeginUser;
646   if (!sgs_dd_data) PetscFunctionReturn(PETSC_SUCCESS);
647   Ceed ceed = sgs_dd_data->op_sgs_apply_ctx->ceed;
648 
649   PetscCallCeed(ceed, CeedVectorDestroy(&sgs_dd_data->sgs_nodal_ceed));
650   PetscCallCeed(ceed, CeedVectorDestroy(&sgs_dd_data->grad_velo_ceed));
651   PetscCall(OperatorApplyContextDestroy(sgs_dd_data->op_nodal_evaluation_ctx));
652   PetscCall(OperatorApplyContextDestroy(sgs_dd_data->op_sgs_apply_ctx));
653   PetscCall(OperatorApplyContextDestroy(sgs_dd_data->op_nodal_dd_inputs_ctx));
654   PetscCall(OperatorApplyContextDestroy(sgs_dd_data->op_nodal_dd_outputs_ctx));
655   PetscCall(DMDestroy(&sgs_dd_data->dm_sgs));
656   PetscCall(DMDestroy(&sgs_dd_data->dm_dd_inputs));
657   PetscCall(DMDestroy(&sgs_dd_data->dm_dd_outputs));
658   if (sgs_dd_data->sgs_nodal_inference_ctx) PetscCall(sgs_dd_data->sgs_nodal_inference_ctx_destroy(sgs_dd_data->sgs_nodal_inference_ctx));
659   PetscCall(PetscFree(sgs_dd_data));
660   PetscFunctionReturn(PETSC_SUCCESS);
661 }
662