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