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