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