Restart netCDF file on checkpoint restart (NGWPC-10459)

include/core/Layer.hpp

@@ -132,7 +132,7 @@ namespace ngen
         std::string unit_name() const;
         virtual std::vector<std::string> required_checkpoint_units() const;
 
-        virtual std::map<std::string, std::string> get_catchment_output_data_for_timestep();
+        virtual const std::map<std::string, std::string>& get_catchment_output_data_for_timestep();
         virtual void set_simulations_output_format(std::vector<std::string> out_formats);
         virtual std::vector<std::string> get_simulations_output_format();
         protected:

include/core/NgenSimulation.hpp

@@ -101,7 +101,7 @@ class NgenSimulation
      * @param merge_all_ranks Whether to also get the units from other MPI ranks. This should only be `true` when calling blocking MPI processes is safe for the program.
      */
     std::vector<std::string> required_checkpoint_units(bool merge_all_ranks) const;
-    void create_netcdf_writer(std::shared_ptr<realization::Formulation_Manager> manager, std::string nc_output_file_name);
+    void create_netcdf_writer(std::shared_ptr<realization::Formulation_Manager> manager, std::string nc_output_file_name, bool create_new_file);
 
 private:
     void advance_models_one_output_step();

include/netcdf/NetCDFFile.hpp

@@ -16,10 +16,17 @@
 #endif
 #include "NetCDFVar.hpp"
 
+
+enum class NetCDFOpenMode {
+    OPEN_READ,
+    OPEN_WRITE,
+    CREATE
+};
+
 class NetCDFFile {
 public:
 
-    NetCDFFile(const std::string& filename, bool write_only, bool is_mpi);
+    NetCDFFile(const std::string& filename, NetCDFOpenMode open_mode, bool is_mpi);
     ~NetCDFFile();
 
     void load_attributes(std::shared_ptr<NetCDFVar> nc_var) ;
@@ -28,6 +35,7 @@ class NetCDFFile {
 
     // Dimension handling
     int add_dimension(const std::string& name, size_t len);
+    int read_dimension_definition(const std::string& name);
     size_t get_dim_size(const std::string& name) const;
     int get_dim_id(const std::string& name) const;
 
@@ -36,6 +44,7 @@ class NetCDFFile {
     std::vector<std::string> list_variables() const;
     void add_ncvar(std::shared_ptr<NetCDFVar> var);
     void add_variable(const std::string& name, nc_type type, const std::vector<int>& dims, const std::vector<std::string>& dim_names);
+    std::shared_ptr<NetCDFVar> read_variable_definition(const std::string& name);
 
     // Write data to a variable
     template<typename T>

include/netcdf/NetCDFManager.hpp

@@ -28,10 +28,10 @@ class NetCDFManager
 {
 public:
      NetCDFManager(std::shared_ptr<realization::Formulation_Manager> manager, 
-        const std::string& output_name, Simulation_Time const& sim_time, int mpi_rank, int mpi_num_procs);
+        const std::string& output_name, Simulation_Time const& sim_time, bool create_new_file, int mpi_rank, int mpi_num_procs);
 
     // Constructor for read-only NetCDF (no MPI needed)
-    NetCDFManager(const std::string& filename, bool read_only);
+    NetCDFManager(const std::string& filename, NetCDFOpenMode open_mode);
 
     // Default constructor for mdframe tests 
     NetCDFManager();
@@ -43,9 +43,6 @@ class NetCDFManager
 
     void gather_all_catchments(const std::vector<int64_t>& catchments_in_proc);
 
-    //set up netcdf dimensions and variables
-    void define_catchment_netcdf_components();
-
     // List variable names
     std::vector<std::string> list_variables() const;
 
@@ -68,21 +65,25 @@ class NetCDFManager
 
     // Add variables to the file (for writing)
     void add_output_variable_data_from_formulation();
+    void read_output_variable_data_from_formulation();
 
     // Add catchment output data to the file (for writing)
-    void write_simulations_response_from_formulation(size_t time_index, std::map<std::string, std::string> catchment_output_values);
+    void write_simulations_response_from_formulation(size_t time_index, const std::map<std::string, std::string>& catchment_output_values);
     void primary_netcdf_writer(size_t time_index, const std::map<int64_t, std::string>& catchment_output_values);
     void secondary_netcdf_worker(const std::map<int64_t, std::string>& catchment_output_values);
 
     ~NetCDFManager();
 
 private:
-    bool read_only_;
+    /* Set up netcdf dimensions and variables.
+       Note: A copy of Simulation_Time is passed because the object is expected to be modified and discarded at the end of the function. */
+    void define_catchment_netcdf_components(Simulation_Time sim_time);
+    void read_catchment_netcdf_components();
+    NetCDFOpenMode open_mode_;
     std::string nc_filename_;
     std::unique_ptr<NetCDFFile> nc_file_;
     std::vector<NetCDFVar> vars_;
     std::shared_ptr<realization::Formulation_Manager> manager_;
-    std::shared_ptr<Simulation_Time> sim_time_;
     size_t num_timesteps_;
     int num_catchments_ = 0;
     std::vector<int64_t> catchments_;

src/NGen.cpp

@@ -714,11 +714,6 @@ int run_ngen(int argc, char* argv[], int mpi_num_procs, int mpi_rank) {
                                                        std::move(nexus_indexes),
                                                        mpi_rank,
                                                        mpi_num_procs);
-    #if NGEN_WITH_NETCDF
-        if (std::find(output_formats.begin(), output_formats.end(), "netcdf") != output_formats.end()){
-            simulation->create_netcdf_writer(manager, "catchment_output");
-        }
-    #endif //NGEN_WITH_NETCDF
     auto time_done_init                             = std::chrono::steady_clock::now();
     std::chrono::duration<double> time_elapsed_init = time_done_init - time_start;
     LOG("[TIMING]: Init: " + std::to_string(time_elapsed_init.count()), LogLevel::INFO);
@@ -732,6 +727,9 @@ int run_ngen(int argc, char* argv[], int mpi_num_procs, int mpi_rank) {
         }
     }
 
+#if NGEN_WITH_NETCDF
+    bool create_new_netcdf_file = true;
+#endif // NGEN_WITH_NETCDF
     { // optionally load a checkpoint if configured
         auto checkpoint_loader = state_saving_config.checkpoint_loader();
         if (checkpoint_loader) {
@@ -740,8 +738,18 @@ int run_ngen(int argc, char* argv[], int mpi_num_procs, int mpi_rank) {
             std::shared_ptr<State_Snapshot_Loader> snapshot_loader
                 = checkpoint_loader->initialize_checkpoint_snapshot(required_units);
             simulation->load_checkpoint(snapshot_loader);
+#if NGEN_WITH_NETCDF
+            create_new_netcdf_file = false;
+#endif //NGEN_WITH_NETCDF
         }
     }
+#if NGEN_WITH_NETCDF
+    if (std::find(output_formats.begin(), output_formats.end(), "netcdf") != output_formats.end()) {
+        if (!create_new_netcdf_file)
+            LOG(LogLevel::INFO, "Attempting to open the output of the prior checkpoint run's netCDF output file. The prior file must be located in the default output location to be correctly read.");
+        simulation->create_netcdf_writer(manager, "catchment_output", create_new_netcdf_file);
+    }
+#endif //NGEN_WITH_NETCDF
 
     if (state_saving_config.has_checkpoint_saver()) {
         int checkpoint_frequency;

src/core/Layer.cpp

@@ -188,7 +188,7 @@ void ngen::Layer::load_hot_start(std::shared_ptr<State_Snapshot_Loader> snapshot
     }
 }
 
-std::map<std::string, std::string> ngen::Layer::get_catchment_output_data_for_timestep(){
+const std::map<std::string, std::string>& ngen::Layer::get_catchment_output_data_for_timestep(){
     return catchment_output_values;
 }
 

src/core/NgenSimulation.cpp

@@ -220,8 +220,8 @@ void NgenSimulation::advance_models_one_output_step()
                     std::vector<std::string> output_formats = layer->get_simulations_output_format();
                     if (std::find(output_formats.begin(), output_formats.end(), "netcdf") != output_formats.end()){
                         produce_netcdf_format_ = true;
-                        std::map<std::string, std::string> catchment_output_vals = layer->get_catchment_output_data_for_timestep();
-                        nc_manager_->write_simulations_response_from_formulation(simulation_step_,catchment_output_vals);
+                        const std::map<std::string, std::string> &catchment_output_vals = layer->get_catchment_output_data_for_timestep();
+                        nc_manager_->write_simulations_response_from_formulation(simulation_step_, catchment_output_vals);
                     }
                 #endif //NGEN_WITH_NETCDF
 
@@ -524,10 +524,10 @@ void NgenSimulation::serialize(Archive& ar, const unsigned int version) {
 #endif //NGEN_WITH_NEXUSES
 }
 
-void NgenSimulation::create_netcdf_writer(std::shared_ptr<realization::Formulation_Manager> manager, std::string nc_output_file_name)
+void NgenSimulation::create_netcdf_writer(std::shared_ptr<realization::Formulation_Manager> manager, std::string nc_output_file_name, bool create_new_file)
 {
 #if NGEN_WITH_NETCDF
-    this->nc_manager_ = std::make_unique<NetCDFManager>(manager, nc_output_file_name, *sim_time_, mpi_rank_, mpi_num_procs_);
+    this->nc_manager_ = std::make_unique<NetCDFManager>(manager, nc_output_file_name, *sim_time_, create_new_file, mpi_rank_, mpi_num_procs_);
 #endif
 }
 

src/forcing/NetCDFPerFeatureDataProvider.cpp

@@ -53,7 +53,7 @@ NetCDFPerFeatureDataProvider::NetCDFPerFeatureDataProvider(std::string input_pat
     //nc_set_chunk_cache(sizep, nelemsp, preemptionp);
 
     //open the file
-    nc_manager = std::make_shared<NetCDFManager>(input_path, true);
+    nc_manager = std::make_shared<NetCDFManager>(input_path, NetCDFOpenMode::OPEN_READ);
     nc_manager->open_file();
 
     //nc_get_chunk_cache(&sizep, &nelemsp, &preemptionp);

src/netcdf/NetCDFFile.cpp

@@ -23,19 +23,27 @@
     } while (0)
 
 
-NetCDFFile::NetCDFFile(const std::string& filename, bool write_only, bool is_mpi)
+NetCDFFile::NetCDFFile(const std::string& filename, NetCDFOpenMode open_mode, bool is_mpi)
     : nc_file_name_(filename), is_mpi_(is_mpi)
 {
-    int mode = NC_NETCDF4;
-    if(write_only){
-        read_only_ = false;
-        NC_CHECK(nc_create(nc_file_name_.c_str(), NC_NETCDF4 | NC_CLOBBER, &ncid_), "Creating NetCDF file failed");
+    switch (open_mode) {
+        case NetCDFOpenMode::OPEN_READ:
+            read_only_ = true;
+            NC_CHECK(nc_open(nc_file_name_.c_str(), NC_NOWRITE, &ncid_), "Opening NetCDF file in read-only mode failed.");
+            break;
+        case NetCDFOpenMode::OPEN_WRITE:
+            read_only_ = false;
+            NC_CHECK(nc_open(nc_file_name_.c_str(), NC_WRITE, &ncid_), "Opening NetCDF file in write mode failed.");
+            break;
+        case NetCDFOpenMode::CREATE:
+            read_only_ = false;
+            NC_CHECK(nc_create(nc_file_name_.c_str(), NC_NETCDF4 | NC_CLOBBER, &ncid_), "Creating NetCDF file failed");
+            break;
+        default:
+            std::string err_msg = "Invalid open mode for NetCDFFile";
+            LOG(LogLevel::FATAL, err_msg);
+            throw std::runtime_error(err_msg);
     }
-    else{
-        read_only_ = true;
-        NC_CHECK(nc_open(nc_file_name_.c_str(), NC_NOWRITE, &ncid_), "Opening NetCDF file failed");
-    }
-//    }
     if(read_only_){
         load_variables(); //load all netcdf data to objects.
     } 
@@ -49,6 +57,13 @@ int NetCDFFile::add_dimension(const std::string& name, size_t len) {
     return dimid;
 }
 
+int NetCDFFile::read_dimension_definition(const std::string& name) {
+    int dimid;
+    NC_CHECK(nc_inq_dimid(ncid_, name.c_str(), &dimid), "NetCDF dimension " + name + " could not be found.");
+    dims_id_map_[name] = dimid;
+    return dimid;
+}
+
 void NetCDFFile::add_variable(const std::string& name, nc_type type, const std::vector<int>& dim_ids, const std::vector<std::string>& dim_names) {
     int varid;
     NC_CHECK(nc_def_var(ncid_, name.c_str(), type, dim_ids.size(), dim_ids.data(), &varid), "NetCDF variable definition failed");
@@ -62,6 +77,29 @@ void NetCDFFile::add_variable(const std::string& name, nc_type type, const std::
     add_ncvar(var);
 }
 
+std::shared_ptr<NetCDFVar> NetCDFFile::read_variable_definition(const std::string& name) {
+    int varid, n_dim_ids;
+    nc_type type;
+    NC_CHECK(nc_inq_varid(ncid_, name.c_str(), &varid), "NetCDF variable " + name + " could not be found.");
+    NC_CHECK(nc_inq_vartype(ncid_, varid, &type), "NetCDF variable " + name + " type read failed.");
+    // get the variable's dimensions
+    std::vector<std::string> dim_names;
+    std::vector<int> dim_ids;
+    NC_CHECK(nc_inq_varndims(ncid_, varid, &n_dim_ids), "NetCDF variable " + name + " associated dimensions count read failed.");
+    dim_ids.resize(n_dim_ids);
+    dim_names.reserve(n_dim_ids);
+    NC_CHECK(nc_inq_vardimid(ncid_, varid, dim_ids.data()), "NetCDF variable " + name + " associated dimensions IDs read failed.");
+    char dim_name[NC_MAX_NAME + 1];
+    for (const int dim_id : dim_ids) {
+        NC_CHECK(nc_inq_dimname(ncid_, dim_id, dim_name), "NetCDF dimension name query failed.");
+        std::string dim_str(dim_name);
+        dim_names.push_back(dim_str);
+    }
+    auto var = std::make_shared<NetCDFVar>(name, type, dim_ids, dim_names, varid, ncid_);
+    add_ncvar(var);
+    return var;
+}
+
 // Get dimension length by name
 size_t NetCDFFile::get_dim_size(const std::string& name) const {
     auto it = dims_id_map_.find(name);