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, §ion)); 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, °ree, 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, §ion)); 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, §ion)); 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