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