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