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
SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData * sgs_dd_train)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
SGS_DD_TrainingSetupDataDestroy(SGS_DD_TrainingSetupData sgs_dd_train_setup_data)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
SGS_DD_TrainingCreateDM(DM dm_source,DM * dm_dd_training,PetscInt degree,PetscInt q_extra,PetscInt * num_components)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
SetupTrainingDataCalculation(Ceed ceed,Honee honee,ProblemData problem,SGS_DD_TrainingSetupData sgs_dd_train_setup_data)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 PetscInt height = 0, dm_field = 0;
98
99 PetscFunctionBeginUser;
100 PetscCall(HoneeGetContainer(honee, SGS_DD_TRAIN_KEY, &sgs_dd_train));
101 PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(sgs_dd_train_setup_data->elem_restr_grid_aniso, &num_comp_grid_aniso));
102
103 PetscCall(DMPlexCeedElemRestrictionCreate(ceed, sgs_dd_train->dm_dd_training, DMLABEL_DEFAULT, DMLABEL_DEFAULT_VALUE, height, dm_field,
104 &elem_restr_sgs_train));
105 PetscCall(GetInverseMultiplicity(ceed, sgs_dd_train->dm_dd_training, DMLABEL_DEFAULT, DMLABEL_DEFAULT_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
SGS_DD_TrainingSetup(Ceed ceed,Honee honee)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 PetscCall(PetscNew(&sgsdd_train_ctx));
179 PetscCall(PetscNew(&sgs_dd_train_setup_data));
180 PetscCall(PetscNew(&sgs_dd_train));
181 *sgs_dd_train = (struct SGS_DD_TrainingData_){
182 .overwrite_training_data = PETSC_TRUE,
183 .write_data_interval = 1,
184 .num_filter_widths = sizeof(sgs_dd_train->filter_widths) / sizeof(sgs_dd_train->filter_widths[0]),
185 };
186 PetscCall(HoneeSetContainer(honee, SGS_DD_TRAIN_KEY, sgs_dd_train, (PetscCtxDestroyFn *)SGS_DD_TrainingDataDestroy));
187 PetscCall(DifferentialFilterSetup(honee, &sgs_dd_train->diff_filter));
188
189 PetscOptionsBegin(honee->comm, NULL, "SGS Data-Driven Training Options", NULL);
190 PetscCall(PetscOptionsInt("-sgs_train_write_data_interval", "Number of timesteps between writing data into database", NULL,
191 sgs_dd_train->write_data_interval, &sgs_dd_train->write_data_interval, NULL));
192 PetscCall(PetscOptionsBool("-sgs_train_overwrite_data", "Overwrite old training data in the database", NULL, sgs_dd_train->overwrite_training_data,
193 &sgs_dd_train->overwrite_training_data, NULL));
194 PetscCall(PetscOptionsRealArray("-sgs_train_filter_width_scales", "Scales of each filter width put into training database", NULL,
195 sgs_dd_train->filter_widths, &sgs_dd_train->num_filter_widths, NULL));
196 PetscOptionsEnd();
197
198 // -- Create DM for storing training data
199 PetscCall(SGS_DD_TrainingCreateDM(honee->dm, &sgs_dd_train->dm_dd_training, honee->app_ctx->degree, honee->app_ctx->q_extra,
200 &sgs_dd_train->num_comp_dd_inputs));
201
202 { // -- Create QFunction Context
203 NewtonianIdealGasContext newt_ctx;
204 PetscCallCeed(ceed, CeedQFunctionContextGetDataRead(problem->apply_vol_ifunction.qfctx, CEED_MEM_HOST, &newt_ctx));
205 sgsdd_train_ctx->newt_ctx = *newt_ctx;
206 PetscCallCeed(ceed, CeedQFunctionContextRestoreDataRead(problem->apply_vol_ifunction.qfctx, &newt_ctx));
207 PetscCallCeed(ceed, CeedQFunctionContextCreate(honee->ceed, &sgs_dd_train_setup_data->sgs_dd_train_qfctx));
208 PetscCallCeed(ceed, CeedQFunctionContextSetData(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, CEED_USE_POINTER,
209 sizeof(*sgsdd_train_ctx), sgsdd_train_ctx));
210 PetscCallCeed(ceed, CeedQFunctionContextSetDataDestroy(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, FreeContextPetsc));
211 }
212
213 { // -- Send training data array info to SmartRedis database
214 PetscMPIInt rank, num_ranks;
215 SmartSimData smartsim;
216 PetscCall(HoneeGetSmartSimData(honee, &smartsim));
217 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank));
218 PetscCallMPI(MPI_Comm_size(honee->comm, &num_ranks));
219
220 {
221 PetscSection global_section;
222 PetscInt num_dofs, num_comps, local_min_max[2] = {0.}, global_min_max[2] = {0.};
223
224 PetscCall(DMGetGlobalSection(sgs_dd_train->dm_dd_training, &global_section));
225 PetscCall(DMGetGlobalVectorInfo(sgs_dd_train->dm_dd_training, &num_dofs, NULL, NULL));
226 PetscCall(PetscSectionGetFieldComponents(global_section, 0, &num_comps));
227 local_min_max[0] = num_dofs;
228 PetscCall(PetscGlobalMinMaxInt(honee->comm, local_min_max, global_min_max));
229
230 sgs_dd_train->training_data_array_dims[0] = global_min_max[0] / num_comps;
231 sgs_dd_train->training_data_array_dims[1] = num_comps;
232 }
233
234 if (rank % smartsim->collocated_database_num_ranks == 0) {
235 { // Communicate info on simulation size
236 const char tensor_name[] = "sizeInfo";
237 size_t array_info_dim = 6;
238 PetscInt64 array_info[6] = {0}, num_features = 6;
239
240 array_info[0] = sgs_dd_train->training_data_array_dims[0];
241 array_info[1] = sgs_dd_train->training_data_array_dims[1];
242 array_info[2] = num_features;
243 array_info[3] = num_ranks;
244 array_info[4] = smartsim->collocated_database_num_ranks;
245 array_info[5] = rank;
246
247 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
248 PetscCallSmartRedis(put_tensor(smartsim->client, tensor_name, strlen(tensor_name), array_info, &array_info_dim, 1, SRTensorTypeInt64,
249 SRMemLayoutContiguous));
250 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name)));
251 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
252 }
253
254 { // Send array that communicates if tensors are overwritten in database
255 const char tensor_name[] = "tensor-ow";
256 PetscInt64 tensor_overwrite[2] = {sgs_dd_train->overwrite_training_data};
257 size_t dim_2[1] = {2};
258
259 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
260 PetscCallSmartRedis(put_tensor(smartsim->client, tensor_name, strlen(tensor_name), tensor_overwrite, dim_2, 1, SRTensorTypeInt64,
261 SRMemLayoutContiguous));
262 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name)));
263 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
264 }
265
266 { // Communicate number of filter widths used
267 const char tensor_name[] = "num_filter_widths";
268 PetscInt64 num_filter_widths = sgs_dd_train->num_filter_widths;
269 size_t dim_2 = 1;
270
271 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
272 PetscCallSmartRedis(put_tensor(smartsim->client, tensor_name, strlen(tensor_name), &num_filter_widths, &dim_2, 1, SRTensorTypeInt64,
273 SRMemLayoutContiguous));
274 PetscCall(SmartRedisVerifyPutTensor(smartsim->client, tensor_name, strlen(tensor_name)));
275 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
276 }
277 }
278 }
279
280 // -- Compute and store anisotropy tensor
281 PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, honee, &sgs_dd_train_setup_data->elem_restr_grid_aniso,
282 &sgs_dd_train_setup_data->grid_aniso_ceed));
283
284 // -- Create Nodal Evaluation Operator
285 PetscCall(SetupTrainingDataCalculation(ceed, honee, problem, sgs_dd_train_setup_data));
286
287 PetscCall(SGS_DD_TrainingSetupDataDestroy(sgs_dd_train_setup_data));
288 PetscFunctionReturn(PETSC_SUCCESS);
289 }
290
TSMonitor_SGS_DD_Training(TS ts,PetscInt step_num,PetscReal solution_time,Vec Q,void * ctx)291 PetscErrorCode TSMonitor_SGS_DD_Training(TS ts, PetscInt step_num, PetscReal solution_time, Vec Q, void *ctx) {
292 Honee honee = (Honee)ctx;
293 Ceed ceed = honee->ceed;
294 SGS_DD_TrainingData sgs_dd_train;
295 SmartSimData smartsim;
296 Vec TrainingData;
297 PetscMPIInt rank;
298
299 PetscFunctionBeginUser;
300 PetscCall(HoneeGetSmartSimData(honee, &smartsim));
301 PetscCall(HoneeGetContainer(honee, SGS_DD_TRAIN_KEY, &sgs_dd_train));
302 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank));
303
304 if (step_num % sgs_dd_train->write_data_interval != 0) PetscFunctionReturn(PETSC_SUCCESS);
305 PetscCall(DMGetGlobalVector(sgs_dd_train->dm_dd_training, &TrainingData));
306
307 for (PetscInt filter_index = 0; filter_index < sgs_dd_train->num_filter_widths; filter_index++) {
308 PetscCall(PetscLogEventBegin(HONEE_TrainDataCompute, 0, 0, 0, 0));
309 { // -- Compute and assemble training data
310 Vec FilteredVelocityGradient, FilteredFields, FilteredFields_loc;
311 PetscMemType filtered_fields_mem_type;
312 CeedVector filtered_fields;
313
314 { // Set filter width for the current solve
315 double filter_width_scaling[3];
316 CeedOperator op_mat;
317 Mat A_mat;
318
319 for (int j = 0; j < 3; j++) filter_width_scaling[j] = sgs_dd_train->filter_widths[filter_index];
320 PetscCall(KSPGetOperators(sgs_dd_train->diff_filter->ksp, &A_mat, NULL));
321 PetscCall(MatCeedGetCeedOperators(A_mat, &op_mat, NULL));
322 PetscCall(CeedOperatorSetContextDouble(op_mat, sgs_dd_train->diff_filter->filter_width_scaling_label, filter_width_scaling));
323 }
324
325 PetscCall(DMGetGlobalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields));
326 PetscCall(DMGetLocalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields_loc));
327
328 PetscCall(DifferentialFilterApply(honee, sgs_dd_train->diff_filter, solution_time, Q, FilteredFields));
329 PetscCall(DMGlobalToLocal(sgs_dd_train->diff_filter->dm_filter, FilteredFields, INSERT_VALUES, FilteredFields_loc));
330
331 PetscCall(DMGetGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient));
332 PetscCall(VelocityGradientProjectionApply(sgs_dd_train->filtered_grad_velo_proj, FilteredFields_loc, FilteredVelocityGradient));
333
334 {
335 CeedOperatorField op_field;
336
337 PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->op_training_data_calc_ctx->op, "q", &op_field));
338 PetscCallCeed(ceed, CeedOperatorFieldGetVector(op_field, &filtered_fields));
339 }
340
341 PetscCall(VecPetscToCeed(FilteredFields_loc, &filtered_fields_mem_type, filtered_fields)); // filtered_fields is an implicit input
342 PetscCall(ApplyCeedOperatorGlobalToGlobal(FilteredVelocityGradient, TrainingData, sgs_dd_train->op_training_data_calc_ctx));
343 PetscCall(VecCeedToPetsc(filtered_fields, filtered_fields_mem_type, FilteredFields_loc));
344
345 PetscCall(DMRestoreGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient));
346 PetscCall(DMRestoreGlobalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields));
347 PetscCall(DMRestoreLocalVector(sgs_dd_train->diff_filter->dm_filter, &FilteredFields_loc));
348 PetscCallCeed(ceed, CeedVectorDestroy(&filtered_fields));
349 }
350 PetscCall(PetscLogEventEnd(HONEE_TrainDataCompute, 0, 0, 0, 0));
351
352 { // -- Send training data to SmartSim
353 char array_key[PETSC_MAX_PATH_LEN];
354 size_t array_key_len;
355
356 if (sgs_dd_train->overwrite_training_data) {
357 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.%" PetscInt_FMT, smartsim->rank_id_name, filter_index));
358 } else {
359 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.%" PetscInt_FMT "%" PetscInt_FMT, smartsim->rank_id_name, step_num, filter_index));
360 }
361 PetscCall(PetscStrlen(array_key, &array_key_len));
362
363 {
364 const PetscScalar *training_data;
365 PetscCall(VecGetArrayRead(TrainingData, &training_data));
366 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Write, 0, 0, 0, 0));
367 PetscCallSmartRedis(put_tensor(smartsim->client, array_key, array_key_len, (void *)training_data, sgs_dd_train->training_data_array_dims, 2,
368 SRTensorTypeDouble, SRMemLayoutContiguous));
369 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Write, 0, 0, 0, 0));
370 PetscCall(VecRestoreArrayRead(TrainingData, &training_data));
371 }
372 }
373 }
374
375 if (rank % smartsim->collocated_database_num_ranks == 0) {
376 const char tensor_name[] = "step";
377 size_t dim_2[1] = {2};
378 PetscInt64 step_array[2] = {step_num, step_num};
379
380 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
381 PetscCallSmartRedis(put_tensor(smartsim->client, tensor_name, strlen(tensor_name), step_array, dim_2, 1, SRTensorTypeInt64,
382 SRMemLayoutContiguous));
383 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
384 }
385
386 PetscCall(DMRestoreGlobalVector(sgs_dd_train->dm_dd_training, &TrainingData));
387 PetscFunctionReturn(PETSC_SUCCESS);
388 }
389
TSPostStep_SGS_DD_Training(TS ts)390 PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) {
391 Honee honee;
392 const char check_run_key[] = "check-run";
393 PetscReal check_run[2] = {1};
394 const size_t check_run_dims[1] = {2};
395 size_t check_run_key_size;
396 SmartSimData smartsim;
397
398 PetscFunctionBeginUser;
399 PetscCall(PetscStrlen(check_run_key, &check_run_key_size));
400 PetscCall(TSGetApplicationContext(ts, &honee));
401 PetscCall(HoneeGetSmartSimData(honee, &smartsim));
402
403 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
404 PetscCallSmartRedis(unpack_tensor(smartsim->client, check_run_key, check_run_key_size, check_run, check_run_dims, 1, SRTensorTypeDouble,
405 SRMemLayoutContiguous));
406 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0));
407 if (check_run[0] == 0) {
408 PetscCall(PetscPrintf(honee->comm, "-- Simulation stopped by 'check-run' tensor in Redis database\n"));
409 PetscCall(TSSetConvergedReason(ts, TS_CONVERGED_USER));
410 }
411 PetscFunctionReturn(PETSC_SUCCESS);
412 }
413