Mesh Refinement (#261)

* [Draft] Mesh Refinement - update test example - add Jupyter notebook for testing - add todo's for SyncRho to fill coarser patches with charge from finer patches Co-authored-by: Remi Lehe <[email protected]> * 1 Level Static MR Working * Generalize prob_relative for MR * Fix RealBox Calcs, Deposit Fine-to-Coarse lvls * Add Masks and Avoid Dbl Counting Prepare for filters, too. * Fix lev>0 ProbDomain In AMReX, even partially covering levels still use the same prob domain as the lower levels * Cleanup: Print Commands * Expanding Beam Example: MR Level 1 - [x] AMReX inputs file with 1 MR level - [ ] Python script with 1 MR level * Docs: Cleaning * Remove Debug Print * Expanding Beam Example: MR Level 1 - [x] AMReX inputs file with 1 MR level - [x] Python script with 1 MR level * Inputs: honor pre-set `amr.max_level` Resetting before `initGrids` does not yet work. * Cleanup: Remove `CheckCharge.ipynb` Used for local tests during development. --------- Co-authored-by: Remi Lehe <[email protected]>
ECP-WarpX · Dec 20, 2023 · 7db7894 · 7db7894
1 parent 7329cad
commit 7db7894
Show file tree

Hide file tree

Showing 14 changed files with 305 additions and 60 deletions.
diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -92,7 +92,7 @@ These boxes also contain a field mesh, if space charge calculations are enabled.
 * ``geometry.dynamic_size`` (``boolean``) optional (default: ``true`` for dynamic)
     Use dynamic (``true``) resizing of the field mesh, via ``geometry.prob_relative``, or static sizing (``false``), via ``geometry.prob_lo``/``geometry.prob_hi``.
 
-* ``geometry.prob_relative`` (positive ``float``, unitless) optional (default: ``3.0``)
+* ``geometry.prob_relative`` (positive ``float`` array with ``amr.max_level`` entries, unitless) optional (default: ``3.0 1.0 1.0 ...``)
     By default, we dynamically extract the minimum and maximum of the particle positions in the beam.
     The field mesh spans, per direction, multiple times the maximum physical extent of beam particles, as given by this factor.
     The beam minimum and maximum extent are symmetrically padded by the mesh.

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -45,13 +45,15 @@ General
 
    .. py:property:: prob_relative
 
+      This is a list with ``amr.max_level`` + 1 entries.
+
       By default, we dynamically extract the minimum and maximum of the particle positions in the beam.
       The field mesh spans, per direction, multiple times the maximum physical extent of beam particles, as given by this factor.
       The beam minimum and maximum extent are symmetrically padded by the mesh.
       For instance, ``1.2`` means the mesh will span 10% above and 10% below the beam;
       ``1.0`` means the beam is exactly covered with the mesh.
 
-      Default: ``3.0``.
+      Default: ``[3.0 1.0 1.0 ...]``.
       When set, turns ``dynamic_size`` to ``True``.
 
    .. py:property:: dynamic_size

diff --git a/examples/cfchannel/run_cfchannel_10nC.py b/examples/cfchannel/run_cfchannel_10nC.py
@@ -15,7 +15,7 @@
 sim.n_cell = [48, 48, 40]  # [72, 72, 64] for increased precision
 sim.particle_shape = 2  # B-spline order
 sim.space_charge = True
-sim.prob_relative = 3.0
+sim.prob_relative = [3.0]
 # sim.diagnostics = False  # benchmarking
 sim.slice_step_diagnostics = True
 

diff --git a/examples/expanding_beam/README.rst b/examples/expanding_beam/README.rst
@@ -14,6 +14,11 @@ expands to twice its original size.  This is tested using the second moments of
 
 In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
 
+This test uses mesh-refinement to solve the space charge force.
+The coarse grid wraps the beam maximum extent by 300%, emulating "open boundary" conditions.
+The refined grid in level 1 spans 110% of the beam maximum extent.
+The grid spacing is adaptively adjusted as the beam evolves.
+
 
 Run
 ---

diff --git a/examples/expanding_beam/input_expanding.in b/examples/expanding_beam/input_expanding.in
@@ -34,5 +34,12 @@ monitor.backend = h5
 algo.particle_shape = 2
 algo.space_charge = true
 
-amr.n_cell = 56 56 48
-geometry.prob_relative = 3.0
+# Space charge solver with one MR level
+amr.max_level = 1
+amr.n_cell = 16 16 20
+geometry.prob_relative = 3.0 1.1
+
+# Space charger solver without MR
+#amr.max_level = 0
+#amr.n_cell = 56 56 48
+#geometry.prob_relative = 3.0
diff --git a/examples/expanding_beam/run_expanding.py b/examples/expanding_beam/run_expanding.py
@@ -9,14 +9,18 @@
 import amrex.space3d as amr
 from impactx import ImpactX, RefPart, distribution, elements
 
+pp_amr = amr.ParmParse("amr")
+pp_amr.add("max_level", 1)
+
 sim = ImpactX()
 
 # set numerical parameters and IO control
-sim.n_cell = [56, 56, 48]
+sim.n_cell = [16, 16, 20]
+# sim.max_level = 1  # TODO: not yet implemented
 sim.particle_shape = 2  # B-spline order
 sim.space_charge = True
 sim.dynamic_size = True
-sim.prob_relative = 3.0
+sim.prob_relative = [3.0, 1.1]
 
 # beam diagnostics
 # sim.diagnostics = False  # benchmarking

diff --git a/src/ImpactX.H b/src/ImpactX.H
@@ -100,11 +100,13 @@ namespace impactx
          */
         void init_warning_logger ();
 
-      private:
         //! Tag cells for refinement.  TagBoxArray tags is built on level lev grids.
-        void ErrorEst (int lev, amrex::TagBoxArray& tags, amrex::Real time,
-                               int ngrow) override;
+        void ErrorEst (int lev,
+                       amrex::TagBoxArray& tags,
+                       amrex::Real time,
+                       int ngrow) override;
 
+      private:
         //! Make a new level from scratch using provided BoxArray and DistributionMapping.
         //! Only used during initialization.
         void MakeNewLevelFromScratch (int lev, amrex::Real time, const amrex::BoxArray& ba,

diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp
@@ -23,10 +23,12 @@
 #include <AMReX.H>
 #include <AMReX_AmrParGDB.H>
 #include <AMReX_BLProfiler.H>
+#include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
 #include <AMReX_Print.H>
 #include <AMReX_Utility.H>
 
+#include <iostream>
 #include <memory>
 
 
@@ -82,7 +84,6 @@ namespace impactx
 
         // init blocks / grids & MultiFabs
         AmrCore::InitFromScratch(0.0);
-        amrex::Print() << "boxArray(0) " << boxArray(0) << std::endl;
 
         // alloc particle containers
         //   the lost particles have an extra runtime attribute: s when it was lost
@@ -98,6 +99,12 @@ namespace impactx
 
         // register shortcut
         m_particle_container->SetLostParticleContainer(m_particles_lost.get());
+
+        // print AMReX grid summary
+        if (amrex::ParallelDescriptor::IOProcessor()) {
+            std::cout << "\nGrids Summary:\n";
+            printGridSummary(std::cout, 0, finestLevel());
+        }
     }
 
     void ImpactX::evolve ()
@@ -194,7 +201,7 @@ namespace impactx
                         m_particle_container->DepositCharge(m_rho, this->refRatio());
 
                         // poisson solve in x,y,z
-                        spacecharge::PoissonSolve(*m_particle_container, m_rho, m_phi);
+                        spacecharge::PoissonSolve(*m_particle_container, m_rho, m_phi, this->ref_ratio);
 
                         // calculate force in x,y,z
                         spacecharge::ForceFromSelfFields(m_space_charge_field,

diff --git a/src/initialization/InitAmrCore.cpp b/src/initialization/InitAmrCore.cpp
@@ -130,7 +130,12 @@ namespace details
         // Periodicity (none)
         amrex::Array<int, AMREX_SPACEDIM> const is_periodic{AMREX_D_DECL(0,0,0)};
 
-        amrex::Geometry const geom(domain, rb, amrex::CoordSys::cartesian, is_periodic);
-        return {geom, amr_info};
+        // Mesh-refinement
+        int max_level = 0;
+        pp_amr.query("max_level", max_level);
+        //   amrex::AmrMesh::InitAmrMesh will query amr.ref_ratio or amr.ref_ratio_vect on its own
+        amrex::Vector<amrex::IntVect> const & ref_ratios = amrex::Vector<amrex::IntVect>();
+
+        return {rb, max_level, n_cell_v, amrex::CoordSys::cartesian, ref_ratios, is_periodic};
     }
 } // namespace impactx::initialization
diff --git a/src/initialization/InitMeshRefinement.cpp b/src/initialization/InitMeshRefinement.cpp
@@ -19,29 +19,144 @@
 #include <AMReX_REAL.H>
 #include <AMReX_Utility.H>
 
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <vector>
 
 
 namespace impactx
 {
-    /** Tag cells for refinement.  TagBoxArray tags is built on level lev grids.
+namespace detail
+{
+    amrex::Vector<amrex::Real>
+    read_mr_prob_relative ()
+    {
+        amrex::ParmParse pp_amr("amr");
+        amrex::ParmParse pp_geometry("geometry");
+
+        int max_level = 0;
+        pp_amr.query("max_level", max_level);
+
+        // The box is expanded beyond the min and max of the particle beam.
+        amrex::Vector<amrex::Real> prob_relative(max_level + 1, 1.0);
+        prob_relative[0] = 3.0;  // top/bottom pad the beam on the lowest level by default by its width
+        pp_geometry.queryarr("prob_relative", prob_relative);
+
+        if (prob_relative[0] < 3.0)
+            ablastr::warn_manager::WMRecordWarning(
+                    "ImpactX::read_mr_prob_relative",
+                    "Dynamic resizing of the mesh uses a geometry.prob_relative "
+                    "padding of less than 3 for level 0. This might result in boundary "
+                    "artifacts for space charge calculation. "
+                    "There is no minimum good value for this parameter, consider "
+                    "doing a convergence test.",
+                    ablastr::warn_manager::WarnPriority::high
+            );
+
+        if (prob_relative[0] < 1.0)
+            throw std::runtime_error("geometry.prob_relative must be >= 1.0 (the beam size) on the coarsest level");
+
+        // check that prob_relative[0] > prob_relative[1] > prob_relative[2] ...
+        amrex::Real last_lev_rel = std::numeric_limits<amrex::Real>::max();
+        for (int lev = 0; lev <= max_level; ++lev) {
+            amrex::Real const prob_relative_lvl = prob_relative[lev];
+            if (prob_relative_lvl <= 0.0)
+                throw std::runtime_error("geometry.prob_relative must be strictly positive for all levels");
+            if (prob_relative_lvl > last_lev_rel)
+                throw std::runtime_error("geometry.prob_relative must be descending over refinement levels");
+
+            last_lev_rel = prob_relative_lvl;
+        }
+
+        return prob_relative;
+    }
+}
+    /** Tag cells on level for refinement.
      *
-     * @todo this function is not (yet) implemented.
+     * @param lev the current level to refine
+     * @param tags the TagBoxArray tags are built on level lev grids
+     * @param time current simulation time, unused
+     * @param ngrow unused legacy argument that is always zero
      */
-    void ImpactX::ErrorEst (int lev, amrex::TagBoxArray& tags, amrex::Real time, int ngrow)
+    void ImpactX::ErrorEst (
+        int lev,
+        amrex::TagBoxArray& tags,
+        [[maybe_unused]] amrex::Real time,
+        [[maybe_unused]] int ngrow
+    )
     {
-        // todo
-        amrex::ignore_unused(lev, tags, time, ngrow);
+        // level zero is of size in meters (per dimension):
+        //   rb_0 = beam_width * prob_relative[lvl=0]
+        // level zero is of size in cells:
+        //   nx_0 = amr.n_cell
+        // level one is of size in meters:
+        //   rb_1 = beam_width * prob_relative[lvl=1]
+        // level one is of size in cells:
+        //   nx_1 = amr.n_cell * rb_1                 / rb_0                 * amr.ref_ratio[lvl=0]
+        //        = amr.n_cell * prob_relative[lvl=1] / prob_relative[lvl=0] * amr.ref_ratio[lvl=0]
+        // level one lowest/highest cell index to refine:
+        //   ref_1_lo = nx_1 / nx_0 / 2
+        //   ref_1_hi = ref_1_lo + nx_1
+        // ...
+        // level n is of size in cells:
+        //   nx_n = amr.n_cell   * rb_n                 / rb_0                         * amr.ref_ratio[lvl=n-1] * amr.ref_ratio[lvl=n-2] * ... * amr.ref_ratio[lvl=0]
+        //        = amr.n_cell   * prob_relative[lvl=n] / prob_relative[lvl=0] * amr.ref_ratio[lvl=n-1] * amr.ref_ratio[lvl=n-2] * ... * amr.ref_ratio[lvl=0]
+        //        = n_cell_{n-1} * prob_relative[lvl=n] / prob_relative[lvl=0] * amr.ref_ratio[lvl=n-1]
+        // level n+1 is of size in cells:
+        //   nx_nup = n_cell_n * prob_relative[lvl=n+1] / prob_relative[lvl=0] * amr.ref_ratio[lvl=n]
+        // level n and n+1 have the following difference in cells, if coarsened to level n
+        //   nx_n_diff = (nx_n - nx_nup / amr.ref_ratio[lvl=n])
+        // level n lowest/highest cell index to refine:
+        //   ref_n_lo = small_end_n + nx_n_diff / 2
+        //   ref_n_hi = big_end_n   - nx_n_diff / 2
+
+        const auto dom = boxArray(lev).minimalBox();  // covered index space of the current level
+        auto const r_cell_n = amrex::RealVect(dom.size());  // on the current level
+        amrex::RealVect const r_ref_ratio = ref_ratio[lev];
+
+        // level width relative to beam width
+        amrex::Vector<amrex::Real> const prob_relative = detail::read_mr_prob_relative();
+
+        amrex::RealVect const n_cell_nup = amrex::RealVect(r_cell_n) * r_ref_ratio
+                                           * prob_relative[lev+1] / prob_relative[0];
+        amrex::RealVect const n_cell_diff = (r_cell_n - n_cell_nup / r_ref_ratio);
+
+        amrex::RealVect const r_fine_tag_lo = amrex::RealVect(dom.smallEnd()) + n_cell_diff / 2.0;
+        amrex::RealVect const r_fine_tag_hi = amrex::RealVect(dom.bigEnd())   - n_cell_diff / 2.0;
+
+        amrex::IntVect const fine_tag_lo = {
+            AMREX_D_DECL((int)std::ceil(r_fine_tag_lo[0]), (int)std::ceil(r_fine_tag_lo[1]), (int)std::ceil(r_fine_tag_lo[2]))
+        };
+        amrex::IntVect const fine_tag_hi = {
+            AMREX_D_DECL((int)std::ceil(r_fine_tag_hi[0]), (int)std::ceil(r_fine_tag_hi[1]), (int)std::ceil(r_fine_tag_hi[2]))
+        };
+
+        amrex::Box const fine_tag = {fine_tag_lo, fine_tag_hi};
+
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
+#endif
+        for (amrex::MFIter mfi(tags); mfi.isValid(); ++mfi)
+        {
+            const amrex::Box& bx = mfi.fabbox();
+            const auto& fab = tags.array(mfi);
+            ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+                amrex::IntVect const idx {AMREX_D_DECL(i, j, k)};
+                if (fine_tag.contains(idx)) {
+                    fab(i,j,k) = amrex::TagBox::SET;
+                }
+            });
+        }
     }
 
     /** Make a new level from scratch using provided BoxArray and DistributionMapping.
      *
      * Only used during initialization.
      */
     void ImpactX::MakeNewLevelFromScratch (int lev, amrex::Real time, const amrex::BoxArray& ba,
-                                          const amrex::DistributionMapping& dm)
+                                           const amrex::DistributionMapping& dm)
     {
         amrex::ignore_unused(time);
 
@@ -148,36 +263,27 @@ namespace impactx
         bool dynamic_size = true;
         pp_geometry.query("dynamic_size", dynamic_size);
 
-        amrex::RealBox rb;
+        amrex::Vector<amrex::RealBox> rb(this->finestLevel() + 1);  // extent per level
         if (dynamic_size)
         {
-            // The box is expanded beyond the min and max of particles.
-            // This controlled by the variable `frac` below.
-            amrex::Real frac = 3.0;
-            pp_geometry.query("prob_relative", frac);
-
-            if (frac < 3.0)
-                ablastr::warn_manager::WMRecordWarning(
-                    "ImpactX::ResizeMesh",
-                    "Dynamic resizing of the mesh uses a geometry.prob_relative "
-                    "with less than 3x the beam size. This might result in boundary "
-                    "artifacts for space charge calculation. "
-                    "There is no minimum good value for this parameter, consider "
-                    "doing a convergence test.",
-                    ablastr::warn_manager::WarnPriority::high
-                );
-
-            if (frac < 1.0)
-                throw std::runtime_error("geometry.prob_relative must be >= 1.0 (the beam size) on the coarsest level");
+            // The coarsest level is expanded (or reduced) relative the min and max of particles.
+            auto const prob_relative = detail::read_mr_prob_relative();
 
+            amrex::Real const frac = prob_relative[0];
             amrex::RealVect const beam_min(x_min, y_min, z_min);
             amrex::RealVect const beam_max(x_max, y_max, z_max);
             amrex::RealVect const beam_width(beam_max - beam_min);
 
             amrex::RealVect const beam_padding = beam_width * (frac - 1.0) / 2.0;
             //                           added to the beam extent --^         ^-- box half above/below the beam
-            rb.setLo(beam_min - beam_padding);
-            rb.setHi(beam_max + beam_padding);
+
+            // In AMReX, all levels have the same problem domain, that of the
+            // coarsest level, even if only partly covered.
+            for (int lev = 0; lev <= this->finestLevel(); ++lev)
+            {
+                rb[lev].setLo(beam_min - beam_padding);
+                rb[lev].setHi(beam_max + beam_padding);
+            }
         }
         else
         {
@@ -188,21 +294,28 @@ namespace impactx
             pp_geometry.getarr("prob_lo", prob_lo);
             pp_geometry.getarr("prob_hi", prob_hi);
 
-            rb = {prob_lo.data(), prob_hi.data()};
+            rb[0] = {prob_lo.data(), prob_hi.data()};
+
+            if (this->max_level > 1)
+                amrex::Abort("Did not implement ResizeMesh for static domains and >1 MR levels.");
         }
 
         // updating geometry.prob_lo/hi for consistency
-        amrex::Vector<amrex::Real> const prob_lo = {rb.lo()[0], rb.lo()[1], rb.lo()[2]};
-        amrex::Vector<amrex::Real> const prob_hi = {rb.hi()[0], rb.hi()[1], rb.hi()[2]};
+        amrex::Vector<amrex::Real> const prob_lo = {rb[0].lo()[0], rb[0].lo()[1], rb[0].lo()[2]};
+        amrex::Vector<amrex::Real> const prob_hi = {rb[0].hi()[0], rb[0].hi()[1], rb[0].hi()[2]};
         pp_geometry.addarr("prob_lo", prob_lo);
         pp_geometry.addarr("prob_hi", prob_hi);
 
         // Resize the domain size
-        amrex::Geometry::ResetDefaultProbDomain(rb);
-        for (int lev = 0; lev <= this->max_level; ++lev) {
+        amrex::Geometry::ResetDefaultProbDomain(rb[0]);
+
+        for (int lev = 0; lev <= this->finestLevel(); ++lev)
+        {
             amrex::Geometry g = Geom(lev);
-            g.ProbDomain(rb);
-            amrex::AmrMesh::SetGeometry(lev, g);
+            g.ProbDomain(rb[lev]);
+            SetGeometry(lev, g);
+
+            m_particle_container->SetParticleGeometry(lev, g);
         }
     }
 } // namespace impactx