xref: /honee/src/smartsim/sgs_dd_training.c (revision 7ff16c02837ccc8cf16a3924fef3b67beddf80c9)
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, &section));
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