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