EB checkpoint files (AMReX-Codes#2897)

* support for loading EB from checkpoint file

* add support for writing chkpt file as well

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Src/EB/AMReX_EB2.H

@@ -128,6 +128,15 @@ void Build (const Geometry& geom,
             bool extend_domain_face = ExtendDomainFace(),
             int num_coarsen_opt = NumCoarsenOpt());
 
+
+void BuildFromChkptFile (std::string const& fname,
+                         const Geometry& geom,
+                         int required_coarsening_level,
+                         int max_coarsening_level,
+                         int ngrow = 4,
+                         bool build_coarse_level_by_coarsening = true,
+                         bool extend_domain_face = ExtendDomainFace());
+
 int maxCoarseningLevel (const Geometry& geom);
 int maxCoarseningLevel (IndexSpace const* ebis, const Geometry& geom);
 

Src/EB/AMReX_EB2.cpp

@@ -11,6 +11,7 @@
 #include <AMReX_EB2_GeometryShop.H>
 #include <AMReX_EB2.H>
 #include <AMReX_EB2_IndexSpace_STL.H>
+#include <AMReX_EB2_IndexSpace_chkpt_file.H>
 #include <AMReX_ParmParse.H>
 #include <AMReX.H>
 #include <algorithm>
@@ -239,6 +240,20 @@ void addFineLevels (int num_new_fine_levels)
     }
 }
 
+void
+BuildFromChkptFile (std::string const& fname,
+                    const Geometry& geom, int required_coarsening_level,
+                    int max_coarsening_level, int ngrow, bool build_coarse_level_by_coarsening,
+                    bool a_extend_domain_face)
+{
+    ChkptFile chkpt_file(fname);
+    IndexSpace::push(new IndexSpaceChkptFile(chkpt_file,
+                     geom, required_coarsening_level,
+                     max_coarsening_level, ngrow,
+                     build_coarse_level_by_coarsening,
+                     a_extend_domain_face));
+}
+
 namespace {
 static int comp_max_crse_level (Box cdomain, const Box& domain)
 {

Src/EB/AMReX_EB2_2D_C.cpp

@@ -391,6 +391,13 @@ void build_cells (Box const& bx, Array4<EBCellFlag> const& cell,
         });
     }
 
+    set_connection_flags(bxg1, cell, fx, fy);
+}
+
+void set_connection_flags (Box const& bxg1,
+                           Array4<EBCellFlag> const& cell,
+                           Array4<Type_t> const& fx, Array4<Type_t> const& fy) noexcept
+{
     // Build neighbors.  By default, all neighbors are already set.
     AMREX_HOST_DEVICE_FOR_3D ( bxg1, i, j, k,
     {

Src/EB/AMReX_EB2_3D_C.cpp

@@ -936,6 +936,14 @@ void build_cells (Box const& bx, Array4<EBCellFlag> const& cell,
         return;
     }
 
+    set_connection_flags(bx, bxg1, cell, ctmp, fx, fy, fz);
+}
+
+void set_connection_flags (Box const& bx,
+                           Box const& bxg1, Array4<EBCellFlag> const& cell,
+                           Array4<EBCellFlag> const& ctmp, Array4<Type_t> const& fx,
+                           Array4<Type_t> const& fy, Array4<Type_t> const& fz) noexcept
+{
     // Build neighbors.  By default all 26 neighbors are already set.
     AMREX_HOST_DEVICE_FOR_3D ( bxg1, i, j, k,
     {

Src/EB/AMReX_EB2_C.H

@@ -36,6 +36,9 @@ void build_cells (Box const& bx, Array4<EBCellFlag> const& cell,
                   Real small_volfrac, Geometry const& geom, bool extend_domain_face,
                   int& nsmallcells, int const nmulticuts) noexcept;
 
+void set_connection_flags(Box const& bxg1, Array4<EBCellFlag> const& cell,
+                          Array4<Type_t> const& fx, Array4<Type_t> const& fy) noexcept;
+
 #elif (AMREX_SPACEDIM == 3)
 
 int build_faces (Box const& bx, Array4<EBCellFlag> const& cell,
@@ -67,6 +70,11 @@ void build_cells (Box const& bx, Array4<EBCellFlag> const& cell,
                   bool extend_domain_face, bool cover_multiple_cuts,
                   int& nsmallcells, int& nmulticuts) noexcept;
 
+void set_connection_flags(Box const& bx, Box const& bxg1,
+                          Array4<EBCellFlag> const& cell, Array4<EBCellFlag> const& ctmp,
+                          Array4<Type_t> const& fx, Array4<Type_t> const& fy,
+                          Array4<Type_t> const& fz) noexcept;
+
 #endif
 
 void intercept_to_edge_centroid (AMREX_D_DECL(Array4<Real> const& excent,

Src/EB/AMReX_EB2_IndexSpace_chkpt_file.H

@@ -0,0 +1,47 @@
+#ifndef AMREX_EB2_INDEXSPACE_CHKPTFILE_H_
+#define AMREX_EB2_INDEXSPACE_CHKPTFILE_H_
+#include <AMReX_Config.H>
+
+#include <AMReX_EB2.H>
+#include <AMReX_EB2_Level_chkpt_file.H>
+
+#include <string>
+
+namespace amrex { namespace EB2 {
+
+class IndexSpaceChkptFile
+    : public IndexSpace
+{
+public:
+
+    IndexSpaceChkptFile (const ChkptFile& chkptfile,
+                         const Geometry& geom, int required_coarsening_level,
+                         int max_coarsening_level, int ngrow,
+                         bool build_coarse_level_by_coarsening,
+                         bool extend_domain_face);
+
+    IndexSpaceChkptFile (IndexSpaceChkptFile const&) = delete;
+    IndexSpaceChkptFile (IndexSpaceChkptFile &&) = delete;
+    void operator= (IndexSpaceChkptFile const&) = delete;
+    void operator= (IndexSpaceChkptFile &&) = delete;
+
+    virtual ~IndexSpaceChkptFile () {}
+
+    virtual const Level& getLevel (const Geometry& geom) const final;
+    virtual const Geometry& getGeometry (const Box& dom) const final;
+    virtual const Box& coarsestDomain () const final {
+        return m_geom.back().Domain();
+    }
+    virtual void addFineLevels (int num_new_fine_levels) final;
+
+private:
+
+    Vector<ChkptFileLevel> m_chkpt_file_level;
+    Vector<Geometry> m_geom;
+    Vector<Box> m_domain;
+    Vector<int> m_ngrow;
+};
+
+}}
+
+#endif

Src/EB/AMReX_EB2_IndexSpace_chkpt_file.cpp

@@ -0,0 +1,86 @@
+#include <AMReX_EB2_IndexSpace_chkpt_file.H>
+
+namespace amrex { namespace EB2 {
+
+IndexSpaceChkptFile::IndexSpaceChkptFile (const ChkptFile& chkpt_file,
+                                          const Geometry& geom, int required_coarsening_level,
+                                          int max_coarsening_level, int ngrow,
+                                          bool build_coarse_level_by_coarsening,
+                                          bool extend_domain_face)
+{
+    Gpu::LaunchSafeGuard lsg(true); // Always use GPU
+
+    // build finest level (i.e., level 0) first
+    AMREX_ALWAYS_ASSERT(required_coarsening_level >= 0 && required_coarsening_level <= 30);
+    max_coarsening_level = std::max(required_coarsening_level,max_coarsening_level);
+    max_coarsening_level = std::min(30,max_coarsening_level);
+
+    int ngrow_finest = std::max(ngrow,0);
+    for (int i = 1; i <= required_coarsening_level; ++i) {
+        ngrow_finest *= 2;
+    }
+
+    m_geom.push_back(geom);
+    m_domain.push_back(geom.Domain());
+    m_ngrow.push_back(ngrow_finest);
+    m_chkpt_file_level.reserve(max_coarsening_level+1);
+    m_chkpt_file_level.emplace_back(this, chkpt_file, geom, EB2::max_grid_size, ngrow_finest,
+            extend_domain_face);
+
+    for (int ilev = 1; ilev <= max_coarsening_level; ++ilev)
+    {
+        bool coarsenable = m_geom.back().Domain().coarsenable(2,2);
+        if (!coarsenable) {
+            if (ilev <= required_coarsening_level) {
+                amrex::Abort("IndexSpaceImp: domain is not coarsenable at level "+std::to_string(ilev));
+            } else {
+                break;
+            }
+        }
+
+        int ng = (ilev > required_coarsening_level) ? 0 : m_ngrow.back()/2;
+
+        Box cdomain = amrex::coarsen(m_geom.back().Domain(),2);
+        Geometry cgeom = amrex::coarsen(m_geom.back(),2);
+        m_chkpt_file_level.emplace_back(this, ilev, EB2::max_grid_size, ng, cgeom, m_chkpt_file_level[ilev-1]);
+        if (!m_chkpt_file_level.back().isOK()) {
+            m_chkpt_file_level.pop_back();
+            if (ilev <= required_coarsening_level) {
+                if (build_coarse_level_by_coarsening) {
+                    amrex::Abort("Failed to build required coarse EB level "+std::to_string(ilev));
+                } else {
+                    amrex::Abort("Chkptfile only stored for finest level. Failed to build "+std::to_string(ilev));
+                }
+            } else {
+                break;
+            }
+        }
+        m_geom.push_back(cgeom);
+        m_domain.push_back(cdomain);
+        m_ngrow.push_back(ng);
+    }
+}
+
+const Level&
+IndexSpaceChkptFile::getLevel (const Geometry& geom) const
+{
+    auto it = std::find(std::begin(m_domain), std::end(m_domain), geom.Domain());
+    int i = std::distance(m_domain.begin(), it);
+    return m_chkpt_file_level[i];
+}
+
+const Geometry&
+IndexSpaceChkptFile::getGeometry (const Box& dom) const
+{
+    auto it = std::find(std::begin(m_domain), std::end(m_domain), dom);
+    int i = std::distance(m_domain.begin(), it);
+    return m_geom[i];
+}
+
+void
+IndexSpaceChkptFile::addFineLevels (int /*num_new_fine_levels*/)
+{
+    amrex::Abort("IndexSpaceChkptFile::addFineLevels: not supported");
+}
+
+}}

Src/EB/AMReX_EB2_Level.H

@@ -60,6 +60,8 @@ public:
     const Geometry& Geom () const noexcept { return m_geom; }
     IndexSpace const* getEBIndexSpace () const noexcept { return m_parent; }
 
+    void write_to_chkpt_file (const std::string& fname, bool extend_domain_face, int max_grid_size) const;
+
 protected:
 
     Level (Level && rhs) = default;

Src/EB/AMReX_EB2_Level.cpp

@@ -1,6 +1,7 @@
 
 #include <AMReX_EB2_Level.H>
 #include <AMReX_IArrayBox.H>
+#include <AMReX_EB_chkpt_file.H>
 #include <algorithm>
 
 namespace amrex { namespace EB2 {
@@ -916,4 +917,14 @@ Level::fillLevelSet (MultiFab& levelset, const Geometry& geom) const
     }
 }
 
+void
+Level::write_to_chkpt_file (const std::string& fname, bool extend_domain_face, int max_grid_size) const
+{
+    ChkptFile chkptFile(fname);
+    chkptFile.write_to_chkpt_file(m_grids, m_covered_grids,
+                                  m_volfrac, m_centroid, m_bndryarea, m_bndrycent,
+                                  m_bndrynorm, m_areafrac, m_facecent, m_edgecent, m_levelset,
+                                  m_geom, m_ngrow, extend_domain_face, max_grid_size);
+}
+
 }}

Src/EB/AMReX_EB2_Level_chkpt_file.H

@@ -0,0 +1,31 @@
+#ifndef AMREX_EB2_LEVEL_CHKPT_FILE_H_
+#define AMREX_EB2_LEVEL_CHKPT_FILE_H_
+#include <AMReX_Config.H>
+
+#include <AMReX_EB2_Level.H>
+#include <AMReX_EB_chkpt_file.H>
+
+namespace amrex { namespace EB2 {
+
+class ChkptFileLevel
+    : public GShopLevel<ChkptFile>
+{
+public:
+
+    ChkptFileLevel (IndexSpace const* is, ChkptFile const& chkpt_file, const Geometry& geom,
+                    int max_grid_size, int ngrow, bool extend_domain_face);
+
+    ChkptFileLevel (IndexSpace const* is, int ilev, int max_grid_size, int ngrow,
+                    const Geometry& geom, ChkptFileLevel& fineLevel);
+
+// for cuda support
+    void define_fine_chkpt_file (ChkptFile const& chkpt_file,
+                                 Geometry const& geom, int max_grid_size, int ngrow,
+                                 bool extend_domain_face);
+
+    void finalize_cell_flags (); //sets the connection flags and adjustments to cellflags
+};
+
+}}
+
+#endif