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_training.h" 5 6 #include <differential_filter.h> 7 #include <navierstokes.h> 8 #include <petscdmplex.h> 9 #include <smartsim-impl.h> 10 11 typedef struct SGS_DD_TrainingData_ *SGS_DD_TrainingData; 12 struct SGS_DD_TrainingData_ { 13 DM dm_dd_training; 14 PetscInt num_comp_dd_inputs, write_data_interval, num_filter_widths; 15 PetscScalar filter_widths[16]; 16 OperatorApplyContext op_training_data_calc_ctx; 17 DiffFilterData diff_filter; 18 NodalProjectionData filtered_grad_velo_proj; 19 size_t training_data_array_dims[2]; 20 PetscBool overwrite_training_data; 21 }; 22 23 #define SGS_DD_TRAIN_KEY "SGS Data Driven Training" 24 25 static PetscErrorCode SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData *sgs_dd_train) { 26 SGS_DD_TrainingData sgs_dd_train_ = *sgs_dd_train; 27 28 PetscFunctionBeginUser; 29 if (!sgs_dd_train_) PetscFunctionReturn(PETSC_SUCCESS); 30 PetscCall(OperatorApplyContextDestroy(sgs_dd_train_->op_training_data_calc_ctx)); 31 PetscCall(NodalProjectionDataDestroy(&sgs_dd_train_->filtered_grad_velo_proj)); 32 PetscCall(DMDestroy(&sgs_dd_train_->dm_dd_training)); 33 PetscCall(DifferentialFilterDataDestroy(&sgs_dd_train_->diff_filter)); 34 PetscCall(PetscFree(sgs_dd_train_)); 35 *sgs_dd_train = NULL; 36 PetscFunctionReturn(PETSC_SUCCESS); 37 } 38 39 typedef struct { 40 CeedElemRestriction elem_restr_grid_aniso; 41 CeedVector grid_aniso_ceed; 42 CeedQFunctionContext sgs_dd_train_qfctx; 43 } *SGS_DD_TrainingSetupData; 44 45 static PetscErrorCode SGS_DD_TrainingSetupDataDestroy(SGS_DD_TrainingSetupData sgs_dd_train_setup_data) { 46 Ceed ceed; 47 48 PetscFunctionBeginUser; 49 PetscCall(CeedElemRestrictionGetCeed(sgs_dd_train_setup_data->elem_restr_grid_aniso, &ceed)); 50 51 PetscCallCeed(ceed, CeedElemRestrictionDestroy(&sgs_dd_train_setup_data->elem_restr_grid_aniso)); 52 PetscCallCeed(ceed, CeedVectorDestroy(&sgs_dd_train_setup_data->grid_aniso_ceed)); 53 PetscCallCeed(ceed, CeedQFunctionContextDestroy(&sgs_dd_train_setup_data->sgs_dd_train_qfctx)); 54 PetscCall(PetscFree(sgs_dd_train_setup_data)); 55 PetscCheck(CeedDestroy(&ceed) == CEED_ERROR_SUCCESS, PETSC_COMM_SELF, PETSC_ERR_LIB, "Destroying Ceed object failed"); 56 PetscFunctionReturn(PETSC_SUCCESS); 57 } 58 59 // @brief Create DM for storing data-drive SGS model inputs 60 static PetscErrorCode SGS_DD_TrainingCreateDM(DM dm_source, DM *dm_dd_training, PetscInt degree, PetscInt q_extra, PetscInt *num_components) { 61 PetscSection section; 62 63 PetscFunctionBeginUser; 64 *num_components = 12; 65 66 PetscCall(DMClone(dm_source, dm_dd_training)); 67 PetscCall(DMSetMatrixPreallocateSkip(*dm_dd_training, PETSC_TRUE)); 68 PetscCall(PetscObjectSetName((PetscObject)*dm_dd_training, "Data-Driven SGS Training Data")); 69 70 PetscCall(DMSetupByOrder_FEM(PETSC_TRUE, PETSC_TRUE, degree, 1, q_extra, 1, num_components, *dm_dd_training)); 71 72 PetscCall(DMGetLocalSection(*dm_dd_training, §ion)); 73 PetscCall(PetscSectionSetFieldName(section, 0, "Data-Driven SGS Training Data")); 74 PetscCall(PetscSectionSetComponentName(section, 0, 0, "SGSInput1")); 75 PetscCall(PetscSectionSetComponentName(section, 0, 1, "SGSInput2")); 76 PetscCall(PetscSectionSetComponentName(section, 0, 2, "SGSInput3")); 77 PetscCall(PetscSectionSetComponentName(section, 0, 3, "SGSInput4")); 78 PetscCall(PetscSectionSetComponentName(section, 0, 4, "SGSInput5")); 79 PetscCall(PetscSectionSetComponentName(section, 0, 5, "SGSInput6")); 80 PetscCall(PetscSectionSetComponentName(section, 0, 6, "FilteredSGSXX")); 81 PetscCall(PetscSectionSetComponentName(section, 0, 7, "FilteredSGSYY")); 82 PetscCall(PetscSectionSetComponentName(section, 0, 8, "FilteredSGSZZ")); 83 PetscCall(PetscSectionSetComponentName(section, 0, 9, "FilteredSGSYZ")); 84 PetscCall(PetscSectionSetComponentName(section, 0, 10, "FilteredSGSXZ")); 85 PetscCall(PetscSectionSetComponentName(section, 0, 11, "FilteredSGSXY")); 86 PetscFunctionReturn(PETSC_SUCCESS); 87 }; 88 89 // @brief Create CeedOperator to calculate training data for data-drive SGS model at nodes 90 static PetscErrorCode SetupTrainingDataCalculation(Ceed ceed, Honee honee, ProblemData problem, SGS_DD_TrainingSetupData sgs_dd_train_setup_data) { 91 SGS_DD_TrainingData sgs_dd_train; 92 CeedQFunction qf_sgs_dd_train; 93 CeedOperator op_sgs_dd_train; 94 CeedInt num_comp_grad_velo, num_comp_grid_aniso; 95 CeedVector inv_multiplicity, filtered_fields; 96 CeedElemRestriction elem_restr_inv_multiplicity, elem_restr_grad_velo, elem_restr_sgs_train; 97 DMLabel domain_label = NULL; 98 PetscInt label_value = 0, height = 0, dm_field = 0; 99 100 PetscFunctionBeginUser; 101 PetscCall(HoneeGetContainer(honee, SGS_DD_TRAIN_KEY, &sgs_dd_train)); 102 PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(sgs_dd_train_setup_data->elem_restr_grid_aniso, &num_comp_grid_aniso)); 103 104 PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_train->dm_dd_training, domain_label, label_value, height, dm_field, &elem_restr_sgs_train)); 105 PetscCall(GetInverseMultiplicity(ceed, sgs_dd_train->dm_dd_training, domain_label, label_value, height, dm_field, PETSC_TRUE, 106 &elem_restr_inv_multiplicity, &inv_multiplicity)); 107 108 CeedElemRestriction elem_restr_filtered_state; 109 CeedInt num_comp_filtered_state; 110 { // -- Setup filtered velocity gradient projection 111 CeedBasis basis_filtered_state; 112 CeedOperatorField op_field; 113 PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->diff_filter->op_rhs_ctx->op, "v0", &op_field)); 114 PetscCallCeed(ceed, CeedOperatorFieldGetData(op_field, NULL, &elem_restr_filtered_state, &basis_filtered_state, NULL)); 115 PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_filtered_state, &num_comp_filtered_state)); 116 PetscCall(VelocityGradientProjectionSetup(ceed, honee, problem, STATEVAR_PRIMITIVE, elem_restr_filtered_state, basis_filtered_state, 117 &sgs_dd_train->filtered_grad_velo_proj)); 118 PetscCallCeed(ceed, CeedBasisDestroy(&basis_filtered_state)); 119 // Get velocity gradient information 120 PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->filtered_grad_velo_proj->l2_rhs_ctx->op, "velocity gradient", &op_field)); 121 PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_grad_velo)); 122 PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_grad_velo, &num_comp_grad_velo)); 123 } 124 125 CeedElemRestriction elem_restr_filtered_velo_prod; 126 CeedInt num_comp_filtered_velo_prod; 127 { // Get filtered velocity product information 128 CeedOperatorField op_field; 129 PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->diff_filter->op_rhs_ctx->op, "v1", &op_field)); 130 PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_filtered_velo_prod)); 131 PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_filtered_velo_prod, &num_comp_filtered_velo_prod)); 132 } 133 134 // -- Create operator for generating training data at nodes 135 // Differential Filter only provides filtered primitive variables 136 PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSGS_DDAnisotropicTrainingDataNodal_Prim, 137 ComputeSGS_DDAnisotropicTrainingDataNodal_Prim_loc, &qf_sgs_dd_train)); 138 139 PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_train, sgs_dd_train_setup_data->sgs_dd_train_qfctx)); 140 PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "q", num_comp_filtered_state, CEED_EVAL_NONE)); 141 PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "velocity product", num_comp_filtered_velo_prod, CEED_EVAL_NONE)); 142 PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "gradient velocity", num_comp_grad_velo, CEED_EVAL_NONE)); 143 PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE)); 144 PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "inverse multiplicity", 1, CEED_EVAL_NONE)); 145 PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_train, "training data", sgs_dd_train->num_comp_dd_inputs, CEED_EVAL_NONE)); 146 147 PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_filtered_state, &filtered_fields, NULL)); 148 PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_train, NULL, NULL, &op_sgs_dd_train)); 149 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "q", elem_restr_filtered_state, CEED_BASIS_NONE, filtered_fields)); 150 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "velocity product", elem_restr_filtered_velo_prod, CEED_BASIS_NONE, filtered_fields)); 151 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "gradient velocity", elem_restr_grad_velo, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE)); 152 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "anisotropy tensor", sgs_dd_train_setup_data->elem_restr_grid_aniso, CEED_BASIS_NONE, 153 sgs_dd_train_setup_data->grid_aniso_ceed)); 154 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_NONE, inv_multiplicity)); 155 PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "training data", elem_restr_sgs_train, CEED_BASIS_NONE, CEED_VECTOR_ACTIVE)); 156 157 PetscCall(OperatorApplyContextCreate(sgs_dd_train->filtered_grad_velo_proj->dm, sgs_dd_train->dm_dd_training, ceed, op_sgs_dd_train, NULL, NULL, 158 NULL, NULL, &sgs_dd_train->op_training_data_calc_ctx)); 159 160 PetscCallCeed(ceed, CeedVectorDestroy(&inv_multiplicity)); 161 PetscCallCeed(ceed, CeedVectorDestroy(&filtered_fields)); 162 PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_inv_multiplicity)); 163 PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_filtered_state)); 164 PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_grad_velo)); 165 PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_filtered_velo_prod)); 166 PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_train)); 167 PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_train)); 168 PetscFunctionReturn(PETSC_SUCCESS); 169 } 170 171 PetscErrorCode SGS_DD_TrainingSetup(Ceed ceed, Honee honee) { 172 SGS_DDTrainingContext sgsdd_train_ctx; 173 SGS_DD_TrainingSetupData sgs_dd_train_setup_data; 174 SGS_DD_TrainingData sgs_dd_train; 175 ProblemData problem = honee->problem_data; 176 177 PetscFunctionBeginUser; 178 179 PetscCall(PetscNew(&sgsdd_train_ctx)); 180 PetscCall(PetscNew(&sgs_dd_train_setup_data)); 181 PetscCall(PetscNew(&sgs_dd_train)); 182 *sgs_dd_train = (struct SGS_DD_TrainingData_){ 183 .overwrite_training_data = PETSC_TRUE, 184 .write_data_interval = 1, 185 .num_filter_widths = sizeof(sgs_dd_train->filter_widths) / sizeof(sgs_dd_train->filter_widths[0]), 186 }; 187 PetscCall(HoneeSetContainer(honee, SGS_DD_TRAIN_KEY, sgs_dd_train, (PetscCtxDestroyFn *)SGS_DD_TrainingDataDestroy)); 188 PetscCall(DifferentialFilterSetup(honee, &sgs_dd_train->diff_filter)); 189 190 PetscOptionsBegin(honee->comm, NULL, "SGS Data-Driven Training Options", NULL); 191 PetscCall(PetscOptionsInt("-sgs_train_write_data_interval", "Number of timesteps between writing data into database", NULL, 192 sgs_dd_train->write_data_interval, &sgs_dd_train->write_data_interval, NULL)); 193 PetscCall(PetscOptionsBool("-sgs_train_overwrite_data", "Overwrite old training data in the database", NULL, sgs_dd_train->overwrite_training_data, 194 &sgs_dd_train->overwrite_training_data, NULL)); 195 PetscCall(PetscOptionsRealArray("-sgs_train_filter_width_scales", "Scales of each filter width put into training database", NULL, 196 sgs_dd_train->filter_widths, &sgs_dd_train->num_filter_widths, NULL)); 197 PetscOptionsEnd(); 198 199 // -- Create DM for storing training data 200 PetscCall(SGS_DD_TrainingCreateDM(honee->dm, &sgs_dd_train->dm_dd_training, honee->app_ctx->degree, honee->app_ctx->q_extra, 201 &sgs_dd_train->num_comp_dd_inputs)); 202 203 { // -- Create QFunction Context 204 NewtonianIdealGasContext newt_ctx; 205 PetscCallCeed(ceed, CeedQFunctionContextGetDataRead(problem->apply_vol_ifunction.qfctx, CEED_MEM_HOST, &newt_ctx)); 206 sgsdd_train_ctx->newt_ctx = *newt_ctx; 207 PetscCallCeed(ceed, CeedQFunctionContextRestoreDataRead(problem->apply_vol_ifunction.qfctx, &newt_ctx)); 208 PetscCallCeed(ceed, CeedQFunctionContextCreate(honee->ceed, &sgs_dd_train_setup_data->sgs_dd_train_qfctx)); 209 PetscCallCeed(ceed, CeedQFunctionContextSetData(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, CEED_USE_POINTER, 210 sizeof(*sgsdd_train_ctx), sgsdd_train_ctx)); 211 PetscCallCeed(ceed, CeedQFunctionContextSetDataDestroy(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, FreeContextPetsc)); 212 } 213 214 { // -- Send training data array info to SmartRedis database 215 PetscMPIInt rank, num_ranks; 216 SmartSimData smartsim; 217 PetscCall(HoneeGetSmartSimData(honee, &smartsim)); 218 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank)); 219 PetscCallMPI(MPI_Comm_size(honee->comm, &num_ranks)); 220 221 { 222 PetscSection global_section; 223 PetscInt num_dofs, num_comps, local_min_max[2] = {0.}, global_min_max[2] = {0.}; 224 225 PetscCall(DMGetGlobalSection(sgs_dd_train->dm_dd_training, &global_section)); 226 PetscCall(DMGetGlobalVectorInfo(sgs_dd_train->dm_dd_training, &num_dofs, NULL, NULL)); 227 PetscCall(PetscSectionGetFieldComponents(global_section, 0, &num_comps)); 228 local_min_max[0] = num_dofs; 229 PetscCall(PetscGlobalMinMaxInt(honee->comm, local_min_max, global_min_max)); 230 231 sgs_dd_train->training_data_array_dims[0] = global_min_max[0] / num_comps; 232 sgs_dd_train->training_data_array_dims[1] = num_comps; 233 } 234 235 if (rank % smartsim->collocated_database_num_ranks == 0) { 236 { // Communicate info on simulation size 237 const char tensor_name[] = "sizeInfo"; 238 size_t array_info_dim = 6; 239 PetscInt64 array_info[6] = {0}, num_features = 6; 240 241 array_info[0] = sgs_dd_train->training_data_array_dims[0]; 242 array_info[1] = sgs_dd_train->training_data_array_dims[1]; 243 array_info[2] = num_features; 244 array_info[3] = num_ranks; 245 array_info[4] = smartsim->collocated_database_num_ranks; 246 array_info[5] = rank; 247 248 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 249 PetscCallSmartRedis( 250 put_tensor(smartsim->client, tensor_name, strlen(tensor_name), array_info, &array_info_dim, 1, SRTensorTypeInt64, SRMemLayoutContiguous)); 251 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name))); 252 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 253 } 254 255 { // Send array that communicates if tensors are overwritten in database 256 const char tensor_name[] = "tensor-ow"; 257 PetscInt64 tensor_overwrite[2] = {sgs_dd_train->overwrite_training_data}; 258 size_t dim_2[1] = {2}; 259 260 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 261 PetscCallSmartRedis( 262 put_tensor(smartsim->client, tensor_name, strlen(tensor_name), tensor_overwrite, dim_2, 1, SRTensorTypeInt64, SRMemLayoutContiguous)); 263 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name))); 264 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 265 } 266 267 { // Communicate number of filter widths used 268 const char tensor_name[] = "num_filter_widths"; 269 PetscInt64 num_filter_widths = sgs_dd_train->num_filter_widths; 270 size_t dim_2 = 1; 271 272 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 273 PetscCallSmartRedis( 274 put_tensor(smartsim->client, tensor_name, strlen(tensor_name), &num_filter_widths, &dim_2, 1, SRTensorTypeInt64, SRMemLayoutContiguous)); 275 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name))); 276 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 277 } 278 } 279 } 280 281 // -- Compute and store anisotropy tensor 282 PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, honee, &sgs_dd_train_setup_data->elem_restr_grid_aniso, 283 &sgs_dd_train_setup_data->grid_aniso_ceed)); 284 285 // -- Create Nodal Evaluation Operator 286 PetscCall(SetupTrainingDataCalculation(ceed, honee, problem, sgs_dd_train_setup_data)); 287 288 PetscCall(SGS_DD_TrainingSetupDataDestroy(sgs_dd_train_setup_data)); 289 PetscFunctionReturn(PETSC_SUCCESS); 290 } 291 292 PetscErrorCode TSMonitor_SGS_DD_Training(TS ts, PetscInt step_num, PetscReal solution_time, Vec Q, void *ctx) { 293 Honee honee = (Honee)ctx; 294 Ceed ceed = honee->ceed; 295 SGS_DD_TrainingData sgs_dd_train; 296 SmartSimData smartsim; 297 Vec TrainingData; 298 PetscMPIInt rank; 299 300 PetscFunctionBeginUser; 301 PetscCall(HoneeGetSmartSimData(honee, &smartsim)); 302 PetscCall(HoneeGetContainer(honee, SGS_DD_TRAIN_KEY, &sgs_dd_train)); 303 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank)); 304 305 if (step_num % sgs_dd_train->write_data_interval != 0) PetscFunctionReturn(PETSC_SUCCESS); 306 PetscCall(DMGetGlobalVector(sgs_dd_train->dm_dd_training, &TrainingData)); 307 308 for (PetscInt filter_index = 0; filter_index < sgs_dd_train->num_filter_widths; filter_index++) { 309 PetscCall(PetscLogEventBegin(HONEE_TrainDataCompute, 0, 0, 0, 0)); 310 { // -- Compute and assemble training data 311 Vec FilteredVelocityGradient, FilteredFields, FilteredFields_loc; 312 PetscMemType filtered_fields_mem_type; 313 CeedVector filtered_fields; 314 315 { // Set filter width for the current solve 316 double filter_width_scaling[3]; 317 CeedOperator op_mat; 318 Mat A_mat; 319 320 for (int j = 0; j < 3; j++) filter_width_scaling[j] = sgs_dd_train->filter_widths[filter_index]; 321 PetscCall(KSPGetOperators(sgs_dd_train->diff_filter->ksp, &A_mat, NULL)); 322 PetscCall(MatCeedGetCeedOperators(A_mat, &op_mat, NULL)); 323 PetscCall(CeedOperatorSetContextDouble(op_mat, sgs_dd_train->diff_filter->filter_width_scaling_label, filter_width_scaling)); 324 } 325 326 PetscCall(DMGetGlobalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields)); 327 PetscCall(DMGetLocalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields_loc)); 328 329 PetscCall(DifferentialFilterApply(honee, sgs_dd_train->diff_filter, solution_time, Q, FilteredFields)); 330 PetscCall(DMGlobalToLocal(sgs_dd_train->diff_filter->dm_filter, FilteredFields, INSERT_VALUES, FilteredFields_loc)); 331 332 PetscCall(DMGetGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient)); 333 PetscCall(VelocityGradientProjectionApply(sgs_dd_train->filtered_grad_velo_proj, FilteredFields_loc, FilteredVelocityGradient)); 334 335 { 336 CeedOperatorField op_field; 337 338 PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->op_training_data_calc_ctx->op, "q", &op_field)); 339 PetscCallCeed(ceed, CeedOperatorFieldGetVector(op_field, &filtered_fields)); 340 } 341 342 PetscCall(VecPetscToCeed(FilteredFields_loc, &filtered_fields_mem_type, filtered_fields)); // filtered_fields is an implicit input 343 PetscCall(ApplyCeedOperatorGlobalToGlobal(FilteredVelocityGradient, TrainingData, sgs_dd_train->op_training_data_calc_ctx)); 344 PetscCall(VecCeedToPetsc(filtered_fields, filtered_fields_mem_type, FilteredFields_loc)); 345 346 PetscCall(DMRestoreGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient)); 347 PetscCall(DMRestoreGlobalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields)); 348 PetscCall(DMRestoreLocalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields_loc)); 349 PetscCallCeed(ceed, CeedVectorDestroy(&filtered_fields)); 350 } 351 PetscCall(PetscLogEventEnd(HONEE_TrainDataCompute, 0, 0, 0, 0)); 352 353 { // -- Send training data to SmartSim 354 char array_key[PETSC_MAX_PATH_LEN]; 355 size_t array_key_len; 356 357 if (sgs_dd_train->overwrite_training_data) { 358 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.%" PetscInt_FMT, smartsim->rank_id_name, filter_index)); 359 } else { 360 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.%" PetscInt_FMT "%" PetscInt_FMT, smartsim->rank_id_name, step_num, filter_index)); 361 } 362 PetscCall(PetscStrlen(array_key, &array_key_len)); 363 364 { 365 const PetscScalar *training_data; 366 PetscCall(VecGetArrayRead(TrainingData, &training_data)); 367 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Write, 0, 0, 0, 0)); 368 PetscCallSmartRedis(put_tensor(smartsim->client, array_key, array_key_len, (void *)training_data, sgs_dd_train->training_data_array_dims, 2, 369 SRTensorTypeDouble, SRMemLayoutContiguous)); 370 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Write, 0, 0, 0, 0)); 371 PetscCall(VecRestoreArrayRead(TrainingData, &training_data)); 372 } 373 } 374 } 375 376 if (rank % smartsim->collocated_database_num_ranks == 0) { 377 const char tensor_name[] = "step"; 378 size_t dim_2[1] = {2}; 379 PetscInt64 step_array[2] = {step_num, step_num}; 380 381 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 382 PetscCallSmartRedis( 383 put_tensor(smartsim->client, tensor_name, strlen(tensor_name), step_array, dim_2, 1, SRTensorTypeInt64, SRMemLayoutContiguous)); 384 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 385 } 386 387 PetscCall(DMRestoreGlobalVector(sgs_dd_train->dm_dd_training, &TrainingData)); 388 PetscFunctionReturn(PETSC_SUCCESS); 389 } 390 391 PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) { 392 Honee honee; 393 const char check_run_key[] = "check-run"; 394 PetscReal check_run[2] = {1}; 395 const size_t check_run_dims[1] = {2}; 396 size_t check_run_key_size; 397 SmartSimData smartsim; 398 399 PetscFunctionBeginUser; 400 PetscCall(PetscStrlen(check_run_key, &check_run_key_size)); 401 PetscCall(TSGetApplicationContext(ts, &honee)); 402 PetscCall(HoneeGetSmartSimData(honee, &smartsim)); 403 404 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 405 PetscCallSmartRedis( 406 unpack_tensor(smartsim->client, check_run_key, check_run_key_size, check_run, check_run_dims, 1, SRTensorTypeDouble, SRMemLayoutContiguous)); 407 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 408 if (check_run[0] == 0) { 409 PetscCall(PetscPrintf(honee->comm, "-- Simulation stopped by 'check-run' tensor in Redis database\n")); 410 PetscCall(TSSetConvergedReason(ts, TS_CONVERGED_USER)); 411 } 412 413 PetscFunctionReturn(PETSC_SUCCESS); 414 } 415