fluids: Add TSPostStep for SGS DD Training

- This will shutdown the simulation if there's a tensor called `check-run` equal to zeros
CEED · May 31, 2023 · 635288d · 635288d
1 parent 58d6d30
commit 635288d
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 0 deletions.
diff --git a/examples/fluids/navierstokes.h b/examples/fluids/navierstokes.h
@@ -475,6 +475,7 @@ PetscErrorCode DifferentialFilter_MMS_ICSetup(ProblemData *problem);
 PetscErrorCode SmartSimSetup(User user);
 PetscErrorCode SGS_DD_TrainingSetup(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem);
 PetscErrorCode TSMonitor_SGS_DD_Training(TS ts, PetscInt step_num, PetscReal solution_time, Vec Q, void *ctx);
+PetscErrorCode TSPostStep_SGS_DD_Training(TS ts);
 PetscErrorCode SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData sgs_dd_train);
 
 #endif  // libceed_fluids_examples_navier_stokes_h
diff --git a/examples/fluids/problems/sgs_dd_training.c b/examples/fluids/problems/sgs_dd_training.c
@@ -442,6 +442,26 @@ PetscErrorCode TSMonitor_SGS_DD_Training(TS ts, PetscInt step_num, PetscReal sol
   PetscFunctionReturn(0);
 }
 
+PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) {
+  User         user;
+  const char   check_run_key[]   = "check-run";
+  PetscReal    check_run[2]      = {0};
+  const size_t check_run_dims[1] = {2}, check_run_key_size = strlen(check_run_key);
+
+  PetscFunctionBeginUser;
+  PetscCall(TSGetApplicationContext(ts, &user));
+  SmartSimData smartsim = user->smartsim;
+
+  SmartRedisCall(
+      unpack_tensor(smartsim->client, check_run_key, check_run_key_size, check_run, check_run_dims, 1, SRTensorTypeDouble, SRMemLayoutContiguous));
+  if (check_run[0] == 0) {
+    PetscCall(PetscPrintf(user->comm, "-- Simulation stopped by 'check-run' tensor in Redis database\n"));
+    PetscCall(TSSetConvergedReason(ts, TS_CONVERGED_USER));
+  }
+
+  PetscFunctionReturn(0);
+}
+
 PetscErrorCode SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData sgs_dd_train) {
   PetscFunctionBeginUser;
   if (!sgs_dd_train) PetscFunctionReturn(0);

diff --git a/examples/fluids/src/setupts.c b/examples/fluids/src/setupts.c
@@ -426,6 +426,7 @@ PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys, Vec *Q
 
   PetscCall(TSCreate(comm, ts));
   PetscCall(TSSetDM(*ts, dm));
+  PetscCall(TSSetApplicationContext(*ts, user));
   if (phys->implicit) {
     PetscCall(TSSetType(*ts, TSBDF));
     if (user->op_ifunction) {

diff --git a/examples/fluids/src/smartsim_weak.c b/examples/fluids/src/smartsim_weak.c
@@ -29,3 +29,9 @@ PetscErrorCode SGS_DD_TrainingDataDestroy(SGS_DD_TrainingData sgs_dd_train) {
   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Warning: SGS_DD_TrainingData struct should not be initialized if SMARTREDIS_DIR isn't set on build..."));
   PetscFunctionReturn(0);
 }
+
+PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) __attribute__((weak));
+PetscErrorCode TSPostStep_SGS_DD_Training(TS ts) {
+  PetscFunctionBeginUser;
+  SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Must build with SMARTREDIS_DIR set to run %s", __func__);
+};