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