Checkpoint tools

amr-wind/utilities/IOManager.H

@@ -45,7 +45,8 @@ public:
     void write_plot_file();
 
     //! Write all necessary fields for restart
-    void write_checkpoint_file(const int start_level = 0);
+    void
+    write_checkpoint_file(const int start_level = 0, const int end_level = -1);
 
     //! Read all necessary fields for a restart
     void read_checkpoint_fields(
@@ -85,8 +86,16 @@ public:
 
     const amrex::Vector<Field*>& plot_fields() const { return m_plt_fields; }
 
+    const amrex::Vector<Field*>& checkpoint_fields() const
+    {
+        return m_chk_fields;
+    }
+
 private:
-    void write_header(const std::string& /*chkname*/, const int start_level);
+    void write_header(
+        const std::string& /*chkname*/,
+        const int start_level,
+        const int end_level);
 
     void write_info_file(const std::string& /*path*/);
 

amr-wind/utilities/IOManager.cpp

@@ -181,7 +181,7 @@ void IOManager::write_plot_file()
 #endif
 }
 
-void IOManager::write_checkpoint_file(const int start_level)
+void IOManager::write_checkpoint_file(const int start_level, int end_level)
 {
     BL_PROFILE("amr-wind::IOManager::write_checkpoint_file");
     const std::string level_prefix = "Level_";
@@ -191,12 +191,16 @@ void IOManager::write_checkpoint_file(const int start_level)
     amrex::Print() << "Writing checkpoint file " << chkname << " at time "
                    << m_sim.time().new_time() << std::endl;
     const auto& mesh = m_sim.mesh();
+    // Modify end_level if need be
+    end_level = (end_level == -1) ? mesh.finestLevel() : end_level;
+    end_level = std::max(end_level, start_level);
+
     amrex::PreBuildDirectorHierarchy(
-        chkname, level_prefix, mesh.finestLevel() + 1 - start_level, true);
-    write_header(chkname, start_level);
+        chkname, level_prefix, end_level + 1 - start_level, true);
+    write_header(chkname, start_level, end_level);
     write_info_file(chkname);
 
-    for (int lev = start_level; lev < mesh.finestLevel() + 1; ++lev) {
+    for (int lev = start_level; lev < end_level + 1; ++lev) {
         for (auto* fld : m_chk_fields) {
             auto& field = *fld;
             amrex::VisMF::Write(
@@ -293,7 +297,8 @@ void IOManager::read_checkpoint_fields(
     }
 }
 
-void IOManager::write_header(const std::string& chkname, const int start_level)
+void IOManager::write_header(
+    const std::string& chkname, const int start_level, const int end_level)
 {
     if (!amrex::ParallelDescriptor::IOProcessor()) {
         return;
@@ -317,7 +322,7 @@ void IOManager::write_header(const std::string& chkname, const int start_level)
     const auto& mesh = m_sim.mesh();
     const auto& time = m_sim.time();
     hdr << "Checkpoint version: 1\n"
-        << mesh.finestLevel() - start_level << "\n"
+        << end_level - start_level << "\n"
         << time.time_index() << "\n"
         << time.new_time() << "\n"
         << time.deltaT() << "\n"
@@ -334,7 +339,7 @@ void IOManager::write_header(const std::string& chkname, const int start_level)
     }
     hdr << "\n";
 
-    for (int lev = start_level; lev < mesh.finestLevel() + 1; ++lev) {
+    for (int lev = start_level; lev < end_level + 1; ++lev) {
         mesh.boxArray(lev).writeOn(hdr);
         hdr << "\n";
     }

tools/utilities/CMakeLists.txt

@@ -1 +1,4 @@
 add_subdirectory(refine-chkpt)
+add_subdirectory(coarsen-chkpt)
+add_subdirectory(CheckpointToCSV)
+add_subdirectory(PlotfileToCSV)

tools/utilities/CheckpointToCSV/CMakeLists.txt

@@ -0,0 +1,17 @@
+set(tool_exe_name CheckpointToCSV)
+
+add_executable(${tool_exe_name})
+target_sources(${tool_exe_name}
+  PRIVATE
+  CheckpointToCSV.cpp
+  CheckpointFileUtil.cpp)
+
+target_include_directories(${tool_exe_name} PRIVATE
+  $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>)
+target_link_libraries(${tool_exe_name} PUBLIC ${amr_wind_lib_name} AMReX-Hydro::amrex_hydro_api)
+set_cuda_build_properties(${tool_exe_name})
+
+install(TARGETS ${tool_exe_name}
+  RUNTIME DESTINATION bin
+  ARCHIVE DESTINATION lib
+  LIBRARY DESTINATION lib)

tools/utilities/CheckpointToCSV/CheckpointFileUtil.H

@@ -0,0 +1,108 @@
+#ifndef CHECKPOINTFILEUTIL_H_
+#define CHECKPOINTFILEUTIL_H_
+#include <AMReX_Config.H>
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_VisMF.H>
+#include <string>
+#include <memory>
+
+using namespace amrex;
+
+// Adapted from amrex/Src/Base/AMReX_PlotFileDataImpl.H
+// not inherited because constructor needed to be rewritten
+
+class CheckpointFileDataImpl
+{
+public:
+    explicit CheckpointFileDataImpl(
+        std::string const& chkptfile_name,
+        Vector<std::string> addl_field_names = {});
+
+    [[nodiscard]] int spaceDim() const noexcept { return m_spacedim; }
+
+    [[nodiscard]] Real time() const noexcept { return m_time; }
+
+    [[nodiscard]] int finestLevel() const noexcept { return m_finest_level; }
+
+    [[nodiscard]] const BoxArray& boxArray(int level) const noexcept
+    {
+        return m_ba[level];
+    }
+
+    [[nodiscard]] const DistributionMapping&
+    DistributionMap(int level) const noexcept
+    {
+        return m_dmap[level];
+    }
+
+    void syncDistributionMap(CheckpointFileDataImpl const& src) noexcept;
+
+    void
+    syncDistributionMap(int level, CheckpointFileDataImpl const& src) noexcept;
+
+    [[nodiscard]] Box probDomain(int level) const noexcept
+    {
+        return m_prob_domain[level];
+    }
+    [[nodiscard]] Array<Real, AMREX_SPACEDIM> probLo() const noexcept
+    {
+        return m_prob_lo;
+    }
+    [[nodiscard]] Array<Real, AMREX_SPACEDIM> probHi() const noexcept
+    {
+        return m_prob_hi;
+    }
+    [[nodiscard]] Array<Real, AMREX_SPACEDIM> cellSize(int level) const noexcept
+    {
+        return m_cell_size[level];
+    }
+
+    [[nodiscard]] const Vector<std::string>& varNames() const noexcept
+    {
+        return m_var_names;
+    }
+
+    [[nodiscard]] const Vector<std::string>& varNamesComponents() const noexcept
+    {
+        return m_var_names_full;
+    }
+
+    [[nodiscard]] int nComp() const noexcept { return m_ncomp; }
+    [[nodiscard]] IntVect nGrowVect(int level) const noexcept
+    {
+        return m_ngrow[level];
+    }
+
+    MultiFab get(int level) noexcept;
+
+private:
+    std::string m_chkptfile_name;
+    int m_ncomp{8};
+    int m_nfields{4};
+    Vector<std::string> m_var_names = {"density", "gp", "p", "velocity"};
+    Vector<std::string> m_var_names_full = {
+        "density", "gpx",       "gpy",       "gpz",
+        "p",       "velocityx", "velocityy", "velocityz"};
+    Vector<int> m_var_ncomp = {1, 3, 1, 3};
+    const int m_spacedim{3};
+    Real m_time;
+    Real m_dt_restart;
+    Real m_dt_nm1;
+    Real m_dt_nm2;
+    int m_finest_level, m_nlevels;
+    Array<Real, AMREX_SPACEDIM> m_prob_lo{{AMREX_D_DECL(0., 0., 0.)}};
+    Array<Real, AMREX_SPACEDIM> m_prob_hi{{AMREX_D_DECL(1., 1., 1.)}};
+    Vector<int> m_ref_ratio;
+    Vector<Box> m_prob_domain;
+    int m_nstep;
+    Vector<Array<Real, AMREX_SPACEDIM>> m_cell_size;
+    Vector<std::string> m_mf_name;
+    Vector<std::unique_ptr<VisMF>> m_vismf;
+    Vector<BoxArray> m_ba;
+    Vector<DistributionMapping> m_dmap;
+    Vector<IntVect> m_ngrow;
+};
+
+#endif

tools/utilities/CheckpointToCSV/CheckpointFileUtil.cpp

@@ -0,0 +1,172 @@
+#include "CheckpointFileUtil.H"
+#include <AMReX_AsyncOut.H>
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_FPC.H>
+#include <AMReX_FabArrayUtility.H>
+#include <AMReX_ParallelDescriptor.H>
+#include <algorithm>
+#include <fstream>
+#include <iomanip>
+
+using namespace amrex;
+
+namespace {
+void GotoNextLine(std::istream& is)
+{
+    constexpr std::streamsize bl_ignore_max{100000};
+    is.ignore(bl_ignore_max, '\n');
+}
+} // namespace
+
+CheckpointFileDataImpl::CheckpointFileDataImpl(
+    std::string const& chkptfile_name, Vector<std::string> addl_field_names)
+    : m_chkptfile_name(chkptfile_name)
+{
+    // Add field names from input
+    const int naf = addl_field_names.size();
+    for (int n = 0; n < naf; ++n) {
+        m_var_names.push_back(addl_field_names[n]);
+        m_var_names_full.push_back(addl_field_names[n]);
+        m_var_ncomp.push_back(1);
+    }
+    // Total number of components, assume addl fields are one component
+    m_ncomp += naf;
+    // Total number of fields
+    m_nfields += naf;
+
+    // Header
+    std::string File(chkptfile_name + "/Header");
+    Vector<char> fileCharPtr;
+    ParallelDescriptor::ReadAndBcastFile(File, fileCharPtr);
+    std::istringstream is(
+        std::string(fileCharPtr.dataPtr()), std::istringstream::in);
+
+    std::string line, word;
+
+    // !! -- Start reading from checkpoint file -- !!
+    // Title line
+    std::getline(is, line);
+
+    // Finest level
+    is >> m_finest_level;
+    GotoNextLine(is);
+
+    // Step count
+    is >> m_nstep;
+    GotoNextLine(is);
+
+    // Current time
+    is >> m_time;
+    GotoNextLine(is);
+
+    // Time step size
+    is >> m_dt_restart;
+    GotoNextLine(is);
+
+    is >> m_dt_nm1;
+    GotoNextLine(is);
+
+    is >> m_dt_nm2;
+    GotoNextLine(is);
+
+    // Low coordinates of domain bounding box
+    std::getline(is, line);
+    {
+        std::istringstream lis(line);
+        int i = 0;
+        while (lis >> word) {
+            m_prob_lo[i++] = std::stod(word);
+        }
+    }
+
+    // High coordinates of domain bounding box
+    std::getline(is, line);
+    {
+        std::istringstream lis(line);
+        int i = 0;
+        while (lis >> word) {
+            m_prob_hi[i++] = std::stod(word);
+        }
+    }
+
+    // Read in box array to get domain box
+    m_nlevels = m_finest_level + 1;
+    m_ba.resize(m_nlevels);
+    m_dmap.resize(m_nlevels);
+    m_ngrow.resize(m_nlevels);
+    m_prob_domain.resize(m_nlevels);
+    m_cell_size.resize(
+        m_nlevels, Array<Real, AMREX_SPACEDIM>{{AMREX_D_DECL(1., 1., 1.)}});
+    for (int lev = 0; lev <= m_finest_level; ++lev) {
+        // read in level 'lev' BoxArray from Header
+        m_ba[lev].readFrom(is);
+        GotoNextLine(is);
+        // get minimal box from box array for prob domain
+        m_prob_domain[lev] = m_ba[lev].minimalBox();
+        // make distribution map
+        m_dmap[lev].define(m_ba[lev], ParallelDescriptor::NProcs());
+        // ngrow is set to 0, don't know how to find it properly
+        m_ngrow[lev] = {0, 0, 0};
+    }
+
+    // Calculate cell size for lowest level
+    for (int idim = 0; idim < m_spacedim; ++idim) {
+        m_cell_size[0][idim] = (m_prob_hi[idim] - m_prob_lo[idim]) /
+                               (m_prob_domain[0].bigEnd(idim) -
+                                m_prob_domain[0].smallEnd(idim) + 1);
+    }
+    // Do the higher levels using refratio, assume 2
+    for (int lev = 1; lev <= m_finest_level; ++lev) {
+        for (int idim = 0; idim < m_spacedim; ++idim) {
+            m_cell_size[lev][idim] = m_cell_size[lev - 1][idim];
+        }
+    }
+
+    AMREX_ASSERT(m_nlevels > 0 && m_nlevels <= 1000);
+}
+
+void CheckpointFileDataImpl::syncDistributionMap(
+    CheckpointFileDataImpl const& src) noexcept
+{
+    int nlevs_min = std::min(m_nlevels, src.m_nlevels);
+    for (int ilev = 0; ilev < nlevs_min; ++ilev) {
+        syncDistributionMap(ilev, src);
+    }
+}
+
+void CheckpointFileDataImpl::syncDistributionMap(
+    int level, CheckpointFileDataImpl const& src) noexcept
+{
+    if (level <= src.finestLevel() &&
+        m_dmap[level].size() == src.DistributionMap(level).size()) {
+        m_dmap[level] = src.DistributionMap(level);
+    }
+}
+
+MultiFab CheckpointFileDataImpl::get(int level) noexcept
+{
+    const std::string level_prefix{"Level_"};
+    MultiFab mf(m_ba[level], m_dmap[level], m_ncomp, m_ngrow[level]);
+
+    // Do checkpoint reading, which is a field at a time
+    for (int lev = 0; lev < m_nlevels; ++lev) {
+        int icomp = 0;
+        for (int nf = 0; nf < m_nfields; ++nf) {
+            MultiFab tmp_mfab(
+                m_ba[level], m_dmap[level], m_var_ncomp[nf], m_ngrow[level]);
+
+            // Read current field into temporary fab
+            amrex::VisMF::Read(
+                tmp_mfab,
+                amrex::MultiFabFileFullPrefix(
+                    lev, m_chkptfile_name, level_prefix, m_var_names[nf]));
+
+            // Copy from tmp fab to bigger fab
+            amrex::MultiFab::Copy(
+                mf, tmp_mfab, 0, icomp, m_var_ncomp[nf], m_ngrow[lev]);
+            // Advance index among total components
+            icomp += m_var_ncomp[nf];
+        }
+    }
+    return mf;
+}