SEAS checkpointing of PETSc TS object

app/CMakeLists.txt

@@ -189,6 +189,8 @@ set(APP_COMMON_SRCS
     localoperator/RateAndStateBase.cpp
     #pc/lspoly.c
     pc/register.cpp
+    common/ts_util.c
+    common/vecnest_util.c
 )
 if(${LAPACK_FOUND})
     list(APPEND APP_COMMON_SRCS

app/common/CmdLine.h

@@ -36,6 +36,11 @@ readFromConfigurationFileAndCmdLine(TableSchema<T> const& schema, argparse::Argu
                   << err << std::endl;
         return std::nullopt;
     }
+    if (!rawCfg.contains("ts_checkpoint")) {
+        // If it doesn't exist, add an empty "ts_checkpoint" section
+        // this allows using default checkpoint options in the tandem app
+        rawCfg.insert_or_assign("ts_checkpoint", toml::array{});
+    }
 
     T cfg;
     try {

app/common/PetscTimeSolver.cpp

@@ -2,12 +2,19 @@
 
 namespace tndm {
 
-PetscTimeSolverBase::PetscTimeSolverBase(MPI_Comm comm) {
+PetscTimeSolverBase::PetscTimeSolverBase(MPI_Comm comm, Config const& cfg) {
     CHKERRTHROW(TSCreate(comm, &ts_));
     CHKERRTHROW(TSSetProblemType(ts_, TS_NONLINEAR));
     CHKERRTHROW(TSSetExactFinalTime(ts_, TS_EXACTFINALTIME_MATCHSTEP));
     CHKERRTHROW(TSSetFromOptions(ts_));
 
+    auto const& cfgcp = cfg.ts_checkpoint_config;
+
+    CHKERRTHROW(
+        ts_checkpoint_configure(ts_, cfgcp.save_directory.c_str(), cfgcp.frequency_step,
+                                cfgcp.frequency_cputime_minutes, cfgcp.frequency_time_physical,
+                                static_cast<int>(cfgcp.storage_type), cfgcp.storage_limited_size));
+
     TSType time_scheme;
     CHKERRTHROW(TSGetType(ts_, &time_scheme));
 

app/common/PetscTimeSolver.h

@@ -3,7 +3,11 @@
 
 #include "common/PetscUtil.h"
 #include "common/PetscVector.h"
-
+#include "tandem/SeasConfig.h"
+extern "C" {
+#include "ts_util.h"
+#include "vecnest_util.h"
+}
 #include <petscsystypes.h>
 #include <petscts.h>
 #include <petscvec.h>
@@ -17,7 +21,7 @@ namespace tndm {
 
 class PetscTimeSolverBase {
 public:
-    PetscTimeSolverBase(MPI_Comm comm);
+    PetscTimeSolverBase(MPI_Comm comm, Config const& cfg);
     ~PetscTimeSolverBase();
 
     std::size_t get_step_number() const;
@@ -33,15 +37,19 @@ class PetscTimeSolverBase {
 template <std::size_t NumStateVecs> class PetscTimeSolver : public PetscTimeSolverBase {
 public:
     template <typename TimeOp>
-    PetscTimeSolver(TimeOp& timeop, std::array<std::unique_ptr<PetscVector>, NumStateVecs> state)
-        : PetscTimeSolverBase(timeop.comm()), state_(std::move(state)) {
+    PetscTimeSolver(TimeOp& timeop, std::array<std::unique_ptr<PetscVector>, NumStateVecs> state,
+                    Config const& cfg)
+        : PetscTimeSolverBase(timeop.comm(), cfg), state_(std::move(state)),
+          ts_checkpoint_load_directory(std::move(cfg.ts_checkpoint_config.load_directory)),
+          comm(timeop.comm()) {
 
         Vec x[NumStateVecs];
         for (std::size_t n = 0; n < NumStateVecs; ++n) {
             x[n] = state_[n]->vec();
         }
         MPI_Comm comm;
         CHKERRTHROW(VecCreateNest(timeop.comm(), NumStateVecs, nullptr, x, &ts_state_));
+        CHKERRTHROW(VecNestUpgradeOperations(ts_state_));
 
         std::apply([&timeop](auto&... x) { timeop.initial_condition((*x)...); }, state_);
 
@@ -51,6 +59,39 @@ template <std::size_t NumStateVecs> class PetscTimeSolver : public PetscTimeSolv
     ~PetscTimeSolver() { VecDestroy(&ts_state_); }
 
     void solve(double upcoming_time) {
+        CHKERRTHROW(TSSetUp(ts_));
+
+        if (ts_checkpoint_load_directory.has_value()) {
+            int rank;
+            MPI_Comm_rank(comm, &rank);
+            std::string sload;
+            if (rank == 0) {
+                sload = ts_checkpoint_load_directory.value();
+                if (std::filesystem::is_regular_file(sload)) {
+                    std::cout << "Retrieving the name of the last checkpoint from " << sload
+                              << std::endl;
+                    std::ifstream file(sload);
+                    if (file.is_open()) {
+                        if (std::getline(file, sload)) {
+                            if (not std::filesystem::is_directory(sload)) {
+                                throw std::runtime_error(
+                                    "The first line of the file does not point to an existing "
+                                    "directory.");
+                            }
+                        } else {
+                            throw std::runtime_error("The file is empty.");
+                        }
+                        file.close();
+                    } else {
+                        throw std::runtime_error("Failed to open the file.");
+                    }
+                }
+            } else {
+                MPI_Bcast(&sload[0], sload.size(), MPI_CHAR, 0, MPI_COMM_WORLD);
+            }
+            const char* loadDirectory = sload.c_str();
+            CHKERRTHROW(ts_checkpoint_restart(ts_, loadDirectory));
+        }
         CHKERRTHROW(TSSetMaxTime(ts_, upcoming_time));
         CHKERRTHROW(TSSolve(ts_, ts_state_));
     }
@@ -108,6 +149,8 @@ template <std::size_t NumStateVecs> class PetscTimeSolver : public PetscTimeSolv
 
     std::array<std::unique_ptr<PetscVector>, NumStateVecs> state_;
     Vec ts_state_ = nullptr;
+    std::optional<std::string> ts_checkpoint_load_directory;
+    MPI_Comm comm;
 };
 
 } // namespace tndm