Add comments

1 files changed, 39 insertions, 10 deletions

diff --git a/Source/WarpXFFT.cpp b/Source/WarpXFFT.cpp
index 8c1f425de..f394d5f64 100644
--- a/Source/WarpXFFT.cpp
+++ b/Source/WarpXFFT.cpp
@@ -10,6 +10,15 @@ constexpr int WarpX::FFTData::N;
 
 namespace {
 
+/** \brief Returns an "owner mask" which 1 for all cells, except
+ *  for the duplicated (physical) cells of a nodal grid.
+ *
+ *  More precisely, for these cells (which are represented on several grids)
+ *  the owner mask is 1 only if these cells are at the lower left end of
+ *  the local grid - or if these cells are at the end of the physical domain
+ *  Therefore, there for these cells, there will be only one grid for
+ *  which the owner mask is non-zero.
+ */
 static iMultiFab
 BuildFFTOwnerMask (const MultiFab& mf, const Geometry& geom)
 {
@@ -23,7 +32,7 @@ BuildFFTOwnerMask (const MultiFab& mf, const Geometry& geom)
     const Box& domain_box = amrex::convert(geom.Domain(), ba.ixType());
 
     AMREX_ASSERT(ba.complementIn(domain_box).isEmpty());
-    
+
 #ifdef _OPENMP
 #pragma omp parallel
 #endif
@@ -33,14 +42,17 @@ BuildFFTOwnerMask (const MultiFab& mf, const Geometry& geom)
         const Box& bx = fab.box();
         Box bx2 = bx;
         for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+            // Detect nodal dimensions
             if (bx2.type(idim) == IndexType::NODE) {
+                // Make sure that this grid does not touch the end of
+                // the physical domain.
                 if (bx2.bigEnd(idim) < domain_box.bigEnd(idim)) {
                     bx2.growHi(idim, -1);
                 }
             }
         }
-
         const BoxList& bl = amrex::boxDiff(bx, bx2);
+        // Set owner mask in these cells
         for (const auto& b : bl) {
             fab.setVal(nonowner, b, 0, 1);
         }
@@ -49,6 +61,14 @@ BuildFFTOwnerMask (const MultiFab& mf, const Geometry& geom)
     return mask;
 }
 
+/** \brief Copy the data from the FFT grid to the regular grid
+ *
+ * Because, for nodal grid, some cells are duplicated on several boxes,
+ * special care has to be taken in order to have consistent values on
+ * each boxes when copying this data. Here this is done by setting a
+ * mask, where, for these duplicated cells, the mask is non-zero on only
+ * one box.
+ */
 static void
 CopyDataFromFFTToValid (MultiFab& mf, const MultiFab& mf_fft, const BoxArray& ba_valid_fft, const Geometry& geom)
 {
@@ -57,6 +77,8 @@ CopyDataFromFFTToValid (MultiFab& mf, const MultiFab& mf_fft, const BoxArray& ba
 
     const iMultiFab& mask = BuildFFTOwnerMask(mftmp, geom);
 
+    // Local copy: whenever an MPI rank owns both the data from the FFT
+    // grid and from the regular grid, for overlapping region, copy it locally
 #ifdef _OPENMP
 #pragma omp parallel
 #endif
@@ -70,7 +92,11 @@ CopyDataFromFFTToValid (MultiFab& mf, const MultiFab& mf_fft, const BoxArray& ba
 
         if (srcbox.contains(bx))
         {
+            // Copy the interior region (without guard cells)
             dstfab.copy(srcfab, bx, 0, bx, 0, 1);
+            // Set the value to 0 whenever the mask is 0
+            // (i.e. for nodal duplicated cells, there is a single box
+            // for which the mask is different than 0)
             amrex_fab_setval_ifnot (BL_TO_FORTRAN_BOX(bx),
                                     BL_TO_FORTRAN_FAB(dstfab),
                                     BL_TO_FORTRAN_ANYD(mask[mfi]),
@@ -78,6 +104,9 @@ CopyDataFromFFTToValid (MultiFab& mf, const MultiFab& mf_fft, const BoxArray& ba
         }
     }
 
+    // Global copy: Get the remaining the data from other procs
+    // Use ParallelAdd instead of ParallelCopy, so that the value from
+    // the cell that has non-zero mask is the one which is retained.
     mf.setVal(0.0, 0);
     mf.ParallelAdd(mftmp);
 }
@@ -203,7 +232,7 @@ WarpX::InitFFTComm (int lev)
     // # of processes in the subcommunicator
     int np_fft = nprocs / ngroups_fft;
     AMREX_ALWAYS_ASSERT_WITH_MESSAGE(np_fft*ngroups_fft == nprocs,
-                                     "Number of processes must be divisilbe by number of FFT groups");
+        "Number of processes must be divisible by number of FFT groups");
 
     int myproc = ParallelDescriptor::MyProc();
     // my color in ngroups_fft subcommunicators.  0 <= color_fft < ngroups_fft
@@ -247,12 +276,15 @@ WarpX::FFTDomainDecompsition (int lev, BoxArray& ba_fft, DistributionMapping& dm
     AMREX_ALWAYS_ASSERT(bl.size() == ngroups_fft);
     const Vector<Box>& bldata = bl.data();
 
-    // This is the domain for the group.
+    // This is the domain for the FFT sub-group (including guard cells)
     domain_fft = amrex::grow(bldata[color_fft[lev]], nguards_fft);
-    // Ask FFTW to chop in z-direction into pieces
+    // Ask FFTW to chop the current FFT sub-group domain in the z-direction
+    // and give a chunk to each MPI rank in the current sub-group.
     int nz_fft, z0_fft;
     warpx_fft_domain_decomp(&nz_fft, &z0_fft, BL_TO_FORTRAN_BOX(domain_fft));
-
+    // Each MPI rank adds a box with its chunk of the FFT grid
+    // (given by the above decomposition) to the list `bx_fft`,
+    // then list is shared among all MPI ranks via AllGather
     Vector<Box> bx_fft;
     if (nz_fft > 0) {
         Box b = domain_fft;
@@ -261,18 +293,15 @@ WarpX::FFTDomainDecompsition (int lev, BoxArray& ba_fft, DistributionMapping& dm
     }
     amrex::AllGatherBoxes(bx_fft);
 
+    // Define the AMReX objects for the FFT grid: BoxArray and DistributionMapping
     ba_fft.define(BoxList(std::move(bx_fft)));
     AMREX_ALWAYS_ASSERT(ba_fft.size() == ParallelDescriptor::NProcs());
-
     Vector<int> pmap(ba_fft.size());
     std::iota(pmap.begin(), pmap.end(), 0);
     dm_fft.define(std::move(pmap));
 
-    //
     // For communication between WarpX normal domain and FFT domain, we need to create a
     // special BoxArray ba_valid
-    //
-
     const Box foobox(-nguards_fft-2, -nguards_fft-2);
 
     BoxList bl_valid; // List of boxes: will be filled by the valid part of the subdomains of ba_fft