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