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 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, §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 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_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 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 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 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 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 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 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 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 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 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 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 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