xref: /honee/src/smartsim/sgs_dd_training.c (revision 80f5d3cbf38b01d9c0f2bbc09c16434641e6b34c)
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, &section));
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     CeedVector    multiplicity;
82     CeedQFunction qf_multiplicity;
83     CeedOperator  op_multiplicity;
84     CeedInt       num_comp_q;
85 
86     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q));
87     PetscCallCeed(ceed, CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &multiplicity, NULL));
88     PetscCallCeed(ceed, CeedElemRestrictionGetMultiplicity(ceed_data->elem_restr_q, multiplicity));
89     PetscCall(DMPlexCeedElemRestrictionCollocatedCreate(ceed, sgs_dd_train->dm_dd_training, domain_label, label_value, height, 1,
90                                                         &elem_restr_inv_multiplicity));
91     PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_inv_multiplicity, &inv_multiplicity, NULL));
92 
93     PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, InverseMultiplicity, InverseMultiplicity_loc, &qf_multiplicity));
94     PetscCallCeed(ceed, CeedQFunctionAddInput(qf_multiplicity, "multiplicity", num_comp_q, CEED_EVAL_NONE));
95     PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_multiplicity, "inverse multiplicity", 1, CEED_EVAL_NONE));
96 
97     PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_multiplicity, NULL, NULL, &op_multiplicity));
98     PetscCallCeed(ceed, CeedOperatorSetName(op_multiplicity, "SGS DD Training Inputs - Create Multiplicity Scaling"));
99     PetscCallCeed(ceed, CeedOperatorSetField(op_multiplicity, "multiplicity", ceed_data->elem_restr_q, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE));
100     PetscCallCeed(
101         ceed, CeedOperatorSetField(op_multiplicity, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE));
102 
103     PetscCallCeed(ceed, CeedOperatorApply(op_multiplicity, multiplicity, inv_multiplicity, CEED_REQUEST_IMMEDIATE));
104 
105     PetscCallCeed(ceed, CeedVectorDestroy(&multiplicity));
106     PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_multiplicity));
107     PetscCallCeed(ceed, CeedOperatorDestroy(&op_multiplicity));
108   }
109 
110   CeedElemRestriction elem_restr_filtered_state;
111   CeedInt             num_comp_filtered_state;
112   {  // -- Setup filtered velocity gradient projection
113     CeedBasis         basis_filtered_state;
114     CeedOperatorField op_field;
115     PetscCallCeed(ceed, CeedOperatorGetFieldByName(user->diff_filter->op_rhs_ctx->op, "v0", &op_field));
116     PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_filtered_state));
117     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_filtered_state, &num_comp_filtered_state));
118     PetscCallCeed(ceed, CeedOperatorFieldGetBasis(op_field, &basis_filtered_state));
119     PetscCall(VelocityGradientProjectionSetup(ceed, user, ceed_data, problem, STATEVAR_PRIMITIVE, elem_restr_filtered_state, basis_filtered_state,
120                                               &sgs_dd_train->filtered_grad_velo_proj));
121     // Get velocity gradient information
122     PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->filtered_grad_velo_proj->l2_rhs_ctx->op, "velocity gradient", &op_field));
123     PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_grad_velo));
124     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_grad_velo, &num_comp_grad_velo));
125   }
126 
127   CeedElemRestriction elem_restr_filtered_velo_prod;
128   CeedInt             num_comp_filtered_velo_prod;
129   {  // Get filtered velocity product information
130     CeedOperatorField op_field;
131     PetscCallCeed(ceed, CeedOperatorGetFieldByName(user->diff_filter->op_rhs_ctx->op, "v1", &op_field));
132     PetscCallCeed(ceed, CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_filtered_velo_prod));
133     PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(elem_restr_filtered_velo_prod, &num_comp_filtered_velo_prod));
134   }
135 
136   // -- Create operator for generating training data at nodes
137   // Differential Filter only provides filtered primitive variables
138   PetscCallCeed(ceed, CeedQFunctionCreateInterior(ceed, 1, ComputeSGS_DDAnisotropicTrainingDataNodal_Prim,
139                                                   ComputeSGS_DDAnisotropicTrainingDataNodal_Prim_loc, &qf_sgs_dd_train));
140 
141   PetscCallCeed(ceed, CeedQFunctionSetContext(qf_sgs_dd_train, sgs_dd_train_setup_data->sgs_dd_train_qfctx));
142   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "q", num_comp_filtered_state, CEED_EVAL_NONE));
143   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "velocity product", num_comp_filtered_velo_prod, CEED_EVAL_NONE));
144   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "gradient velocity", num_comp_grad_velo, CEED_EVAL_NONE));
145   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE));
146   PetscCallCeed(ceed, CeedQFunctionAddInput(qf_sgs_dd_train, "inverse multiplicity", 1, CEED_EVAL_NONE));
147   PetscCallCeed(ceed, CeedQFunctionAddOutput(qf_sgs_dd_train, "training data", sgs_dd_train->num_comp_dd_inputs, CEED_EVAL_NONE));
148 
149   PetscCallCeed(ceed, CeedElemRestrictionCreateVector(elem_restr_filtered_state, &filtered_fields, NULL));
150   PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_sgs_dd_train, NULL, NULL, &op_sgs_dd_train));
151   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "q", elem_restr_filtered_state, CEED_BASIS_COLLOCATED, filtered_fields));
152   PetscCallCeed(ceed,
153                 CeedOperatorSetField(op_sgs_dd_train, "velocity product", elem_restr_filtered_velo_prod, CEED_BASIS_COLLOCATED, filtered_fields));
154   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "gradient velocity", elem_restr_grad_velo, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE));
155   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "anisotropy tensor", sgs_dd_train_setup_data->elem_restr_grid_aniso,
156                                            CEED_BASIS_COLLOCATED, sgs_dd_train_setup_data->grid_aniso_ceed));
157   PetscCallCeed(ceed,
158                 CeedOperatorSetField(op_sgs_dd_train, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_COLLOCATED, inv_multiplicity));
159   PetscCallCeed(ceed, CeedOperatorSetField(op_sgs_dd_train, "training data", elem_restr_sgs_train, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE));
160 
161   PetscCall(OperatorApplyContextCreate(sgs_dd_train->filtered_grad_velo_proj->dm, sgs_dd_train->dm_dd_training, ceed, op_sgs_dd_train, NULL, NULL,
162                                        NULL, NULL, &sgs_dd_train->op_training_data_calc_ctx));
163 
164   PetscCallCeed(ceed, CeedVectorDestroy(&inv_multiplicity));
165   PetscCallCeed(ceed, CeedVectorDestroy(&filtered_fields));
166   PetscCallCeed(ceed, CeedElemRestrictionDestroy(&elem_restr_inv_multiplicity));
167   PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_sgs_dd_train));
168   PetscCallCeed(ceed, CeedOperatorDestroy(&op_sgs_dd_train));
169   PetscFunctionReturn(PETSC_SUCCESS);
170 }
171 
172 PetscErrorCode SGS_DD_TrainingSetup(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem) {
173   SGS_DDTrainingContext    sgsdd_train_qfctx;
174   SGS_DD_TrainingSetupData sgs_dd_train_setup_data;
175 
176   PetscFunctionBeginUser;
177   if (!user->diff_filter) PetscCall(DifferentialFilterSetup(ceed, user, ceed_data, problem));
178   if (!user->smartsim) PetscCall(SmartSimSetup(user));
179 
180   PetscCall(PetscNew(&sgsdd_train_qfctx));
181   PetscCall(PetscNew(&sgs_dd_train_setup_data));
182   PetscCall(PetscNew(&user->sgs_dd_train));
183   SGS_DD_TrainingData sgs_dd_train = user->sgs_dd_train;
184 
185   sgs_dd_train->overwrite_training_data = PETSC_TRUE;
186   sgs_dd_train->write_data_interval     = 1;
187   PetscOptionsBegin(user->comm, NULL, "SGS Data-Driven Training Options", NULL);
188   PetscCall(PetscOptionsInt("-sgs_train_write_data_interval", "Number of timesteps between writing data into database", NULL,
189                             sgs_dd_train->write_data_interval, &sgs_dd_train->write_data_interval, NULL));
190   PetscCall(PetscOptionsBool("-sgs_train_overwrite_data", "Overwrite old training data in the database", NULL, sgs_dd_train->overwrite_training_data,
191                              &sgs_dd_train->overwrite_training_data, NULL));
192   PetscOptionsEnd();
193 
194   // -- Create DM for storing training data
195   PetscCall(SGS_DD_TrainingCreateDM(user->dm, &sgs_dd_train->dm_dd_training, user->app_ctx->degree, user->app_ctx->q_extra,
196                                     &sgs_dd_train->num_comp_dd_inputs));
197 
198   {  // -- Create QFunction Context
199     NewtonianIdealGasContext gas;
200     PetscCallCeed(ceed, CeedQFunctionContextGetDataRead(problem->apply_vol_ifunction.qfunction_context, CEED_MEM_HOST, &gas));
201     sgsdd_train_qfctx->gas = *gas;
202     PetscCallCeed(ceed, CeedQFunctionContextRestoreDataRead(problem->apply_vol_ifunction.qfunction_context, &gas));
203     PetscCallCeed(ceed, CeedQFunctionContextCreate(user->ceed, &sgs_dd_train_setup_data->sgs_dd_train_qfctx));
204     PetscCallCeed(ceed, CeedQFunctionContextSetData(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, CEED_USE_POINTER,
205                                                     sizeof(*sgsdd_train_qfctx), sgsdd_train_qfctx));
206     PetscCallCeed(ceed, CeedQFunctionContextSetDataDestroy(sgs_dd_train_setup_data->sgs_dd_train_qfctx, CEED_MEM_HOST, FreeContextPetsc));
207   }
208 
209   {  // -- Send training data array info to SmartRedis database
210     PetscMPIInt  rank, num_ranks;
211     SmartSimData smartsim = user->smartsim;
212     PetscCallMPI(MPI_Comm_rank(user->comm, &rank));
213     PetscCallMPI(MPI_Comm_size(user->comm, &num_ranks));
214 
215     {
216       PetscSection global_section;
217       PetscInt     num_dofs, num_comps;
218       PetscCall(DMGetGlobalSection(sgs_dd_train->dm_dd_training, &global_section));
219       PetscCall(DMGetGlobalVectorInfo(sgs_dd_train->dm_dd_training, &num_dofs, NULL, NULL));
220       PetscCall(PetscSectionGetFieldComponents(global_section, 0, &num_comps));
221       sgs_dd_train->training_data_array_dims[0] = num_dofs / num_comps;
222       sgs_dd_train->training_data_array_dims[1] = num_comps;
223     }
224 
225     if (rank % smartsim->collocated_database_num_ranks == 0) {
226       size_t   array_info_dim = 6;
227       PetscInt array_info[6] = {0}, num_features = 6;
228 
229       array_info[0] = sgs_dd_train->training_data_array_dims[0];
230       array_info[1] = sgs_dd_train->training_data_array_dims[1];
231       array_info[2] = num_features;
232       array_info[3] = num_ranks;
233       array_info[4] = smartsim->collocated_database_num_ranks;
234       array_info[5] = rank;
235 
236       PetscSmartRedisCall(put_tensor(smartsim->client, "sizeInfo", 8, array_info, &array_info_dim, 1, SRTensorTypeInt32, SRMemLayoutContiguous));
237       PetscCall(SmartRedisVerifyPutTensor(smartsim->client, "sizeInfo", 8));
238 
239       // -- Send array that communicates if tensors are overwritten in database
240       PetscInt tensor_overwrite[2] = {sgs_dd_train->overwrite_training_data};
241       size_t   dim_2[1]            = {2};
242       PetscSmartRedisCall(put_tensor(smartsim->client, "tensor-ow", 9, tensor_overwrite, dim_2, 1, SRTensorTypeInt32, SRMemLayoutContiguous));
243       PetscCall(SmartRedisVerifyPutTensor(smartsim->client, "tensor-ow", 9));
244     }
245   }
246 
247   // -- Compute and store anisotropy tensor
248   PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, user, ceed_data, &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, user, ceed_data, 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   User                user         = (User)ctx;
260   Ceed                ceed         = user->ceed;
261   SGS_DD_TrainingData sgs_dd_train = user->sgs_dd_train;
262   SmartSimData        smartsim     = user->smartsim;
263   Vec                 TrainingData;
264 
265   PetscFunctionBeginUser;
266   if (step_num % sgs_dd_train->write_data_interval != 0) PetscFunctionReturn(PETSC_SUCCESS);
267   PetscCall(DMGetGlobalVector(sgs_dd_train->dm_dd_training, &TrainingData));
268 
269   {  // -- Compute and assemble training data
270     Vec          FilteredVelocityGradient, FilteredFields, FilteredFields_loc;
271     PetscMemType filtered_fields_mem_type;
272     CeedVector   filtered_fields;
273 
274     PetscCall(DMGetGlobalVector(user->diff_filter->dm_filter, &FilteredFields));
275     PetscCall(DMGetLocalVector(user->diff_filter->dm_filter, &FilteredFields_loc));
276 
277     PetscCall(DifferentialFilterApply(user, solution_time, Q, FilteredFields));
278     PetscCall(DMGlobalToLocal(user->diff_filter->dm_filter, FilteredFields, INSERT_VALUES, FilteredFields_loc));
279 
280     PetscCall(DMGetGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient));
281     PetscCall(VelocityGradientProjectionApply(sgs_dd_train->filtered_grad_velo_proj, FilteredFields_loc, FilteredVelocityGradient));
282 
283     {
284       CeedOperatorField op_field;
285       PetscCallCeed(ceed, CeedOperatorGetFieldByName(sgs_dd_train->op_training_data_calc_ctx->op, "q", &op_field));
286       PetscCallCeed(ceed, CeedOperatorFieldGetVector(op_field, &filtered_fields));
287     }
288     PetscCall(VecP2C(FilteredFields_loc, &filtered_fields_mem_type, filtered_fields));  // filtered_fields is an implicit input
289 
290     PetscCall(ApplyCeedOperatorGlobalToGlobal(FilteredVelocityGradient, TrainingData, sgs_dd_train->op_training_data_calc_ctx));
291 
292     PetscCall(VecC2P(filtered_fields, filtered_fields_mem_type, FilteredFields_loc));
293 
294     PetscCall(DMRestoreGlobalVector(sgs_dd_train->filtered_grad_velo_proj->dm, &FilteredVelocityGradient));
295     PetscCall(DMRestoreGlobalVector(user->diff_filter->dm_filter, &FilteredFields));
296     PetscCall(DMRestoreLocalVector(user->diff_filter->dm_filter, &FilteredFields_loc));
297   }
298 
299   {  // -- Send training data to SmartSim
300     char        array_key[PETSC_MAX_PATH_LEN];
301     size_t      array_key_len;
302     PetscMPIInt rank;
303 
304     PetscCallMPI(MPI_Comm_rank(user->comm, &rank));
305 
306     if (sgs_dd_train->overwrite_training_data) {
307       PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s", smartsim->rank_id_name));
308     } else {
309       PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.%" PetscInt_FMT, smartsim->rank_id_name, step_num));
310     }
311     PetscCall(PetscStrlen(array_key, &array_key_len));
312 
313     {
314       const PetscScalar *training_data;
315       PetscCall(VecGetArrayRead(TrainingData, &training_data));
316       PetscSmartRedisCall(put_tensor(smartsim->client, array_key, array_key_len, (void *)training_data, sgs_dd_train->training_data_array_dims, 2,
317                                      SRTensorTypeDouble, SRMemLayoutContiguous));
318       PetscCall(VecRestoreArrayRead(TrainingData, &training_data));
319     }
320     PetscCall(SmartRedisVerifyPutTensor(smartsim->client, array_key, array_key_len));
321 
322     if (rank % smartsim->collocated_database_num_ranks == 0) {
323       size_t   dim_2[1]      = {2};
324       PetscInt step_array[2] = {step_num, step_num};
325       PetscSmartRedisCall(put_tensor(smartsim->client, "step", 4, step_array, dim_2, 1, SRTensorTypeInt32, SRMemLayoutContiguous));
326     }
327   }
328 
329   PetscCall(DMRestoreGlobalVector(user->sgs_dd_train->dm_dd_training, &TrainingData));
330   PetscFunctionReturn(PETSC_SUCCESS);
331 }
332 
333 PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) {
334   User         user;
335   const char   check_run_key[]   = "check-run";
336   PetscReal    check_run[2]      = {1};
337   const size_t check_run_dims[1] = {2};
338   size_t       check_run_key_size;
339 
340   PetscFunctionBeginUser;
341   PetscCall(PetscStrlen(check_run_key, &check_run_key_size));
342   PetscCall(TSGetApplicationContext(ts, &user));
343   SmartSimData smartsim = user->smartsim;
344 
345   PetscSmartRedisCall(
346       unpack_tensor(smartsim->client, check_run_key, check_run_key_size, check_run, check_run_dims, 1, SRTensorTypeDouble, SRMemLayoutContiguous));
347   if (check_run[0] == 0) {
348     PetscCall(PetscPrintf(user->comm, "-- Simulation stopped by 'check-run' tensor in Redis database\n"));
349     PetscCall(TSSetConvergedReason(ts, TS_CONVERGED_USER));
350   }
351 
352   PetscFunctionReturn(PETSC_SUCCESS);
353 }
354 
355 PetscErrorCode SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData sgs_dd_train) {
356   PetscFunctionBeginUser;
357   if (!sgs_dd_train) PetscFunctionReturn(PETSC_SUCCESS);
358 
359   PetscCall(OperatorApplyContextDestroy(sgs_dd_train->op_training_data_calc_ctx));
360   PetscCall(NodalProjectionDataDestroy(sgs_dd_train->filtered_grad_velo_proj));
361   PetscCall(DMDestroy(&sgs_dd_train->dm_dd_training));
362   PetscCall(PetscFree(sgs_dd_train));
363 
364   PetscFunctionReturn(PETSC_SUCCESS);
365 }
366