New function for domain decomposition

Add a new function that decomposition a Box into BoxArray. The number of
Boxes in the returned BoxArray matches the given nboxes argument, unless the
domain is too small. We aim to decompose the domain into subdomains that are
as cubic as possible, even if this results in Boxes with odd numbers of
cells. Thus, this function is generally not suited for applications with
multiple AMR levels or for multigrid solvers. However, it could be useful
for single-level applications that prefer exactly one Box per process.

Src/Base/AMReX_BoxArray.H

@@ -53,6 +53,24 @@ namespace amrex
     //! Note that two BoxArrays that match are not necessarily equal.
     [[nodiscard]] bool match (const BoxArray& x, const BoxArray& y);
 
+    /**
+     * \breif Decompose domain box into BoxArray
+     *
+     * The returned BoxArray has nboxes Boxes, unless the the domain is too
+     * small. We aim to decompose the domain into subdomains that are as
+     * cubic as possible, even if this results in Boxes with odd numbers of
+     * cells. Thus, this function is generally not suited for applications
+     * with multiple AMR levels or for multigrid solvers.
+     *
+     * \param domain Domain Box
+     * \param nboxes the target number of Boxes
+     * \param decomp controls whether domain decomposition should be done in
+     *               that direction.
+     */
+    [[nodiscard]] BoxArray decompose (Box const& domain, int nboxes,
+                                      Array<bool,AMREX_SPACEDIM> const& decomp
+                                      = {AMREX_D_DECL(true,true,true)});
+
 struct BARef
 {
     BARef ();

Src/Base/AMReX_BoxArray.cpp

@@ -12,6 +12,9 @@
 
 #include <AMReX_OpenMP.H>
 
+#include <algorithm>
+#include <cstdlib>
+#include <functional>
 #include <iostream>
 
 namespace amrex {
@@ -1887,6 +1890,170 @@ bool match (const BoxArray& x, const BoxArray& y)
     }
 }
 
+BoxArray decompose (Box const& domain, int nboxes,
+                    Array<bool,AMREX_SPACEDIM> const& decomp)
+{
+    int ndecomp = std::count(decomp.begin(), decomp.end(), true);
+
+    if (nboxes <= 1 || ndecomp == 0) {
+        return BoxArray(domain);
+    }
+
+    Box const& ccdomain = amrex::enclosedCells(domain);
+    IntVect const& ncells = ccdomain.length();
+    IntVect nprocs(1);
+
+    if (ndecomp == 1) {
+        for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+            if (decomp[idim]) {
+                nprocs[idim] = nboxes;
+            }
+        }
+    } else {
+        // Factorization of nboxes
+        Vector<int> factors;
+        {
+            int x = 2;
+            int n = nboxes;
+            while (x*x <= n) {
+                std::div_t dv = std::div(n, x);
+                if (dv.rem == 0) {
+                    factors.push_back(x);
+                    n = dv.quot;
+                } else {
+                    ++x;
+                }
+            }
+            if (n != 1) {
+                factors.push_back(n);
+            }
+            AMREX_ALWAYS_ASSERT(nboxes == std::accumulate(factors.begin(), factors.end(),
+                                                          1, std::multiplies<int>()));
+        }
+
+        struct ProcDim
+        {
+            int nproc;
+            int idim;
+            Vector<int> procs;
+            ProcDim (int np, int dim) : nproc(np), idim(dim) {}
+        };
+
+        Vector<ProcDim> procdim;
+        procdim.reserve(AMREX_SPACEDIM);
+
+        Array<Long,AMREX_SPACEDIM> nblocks;
+
+        for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+            if (decomp[idim]) {
+                nblocks[idim] = ncells[idim];
+                procdim.emplace_back(1,idim);
+            } else {
+                nblocks[idim] = 0;  // This dimension will not be decomposed.
+            }
+        }
+
+        auto comp = [&] (ProcDim const& a, ProcDim const& b) {
+                          if (nblocks[a.idim]*b.nproc <
+                              nblocks[b.idim]*a.nproc) {
+                              return true;
+                          } else if (nblocks[a.idim]*b.nproc >
+                                     nblocks[b.idim]*a.nproc) {
+                              return false;
+                          } else {
+                              return a.procs.size() > b.procs.size();
+                          }
+                      };
+
+        int nprocs_tot = 1;
+        while (!factors.empty()) {
+            std::sort(procdim.begin(), procdim.end(), comp);
+            auto f = factors.back();
+            factors.pop_back();
+            procdim.back().nproc *= f;
+            procdim.back().procs.push_back(f);
+            nprocs_tot *= f;
+            if (nprocs_tot == nboxes) {
+                break;
+            }
+        }
+
+        // swap to see if the decomposition can be improved.
+        while (true)
+        {
+            std::sort(procdim.begin(), procdim.end(), comp);
+            auto& light = procdim.front();
+            auto& heavy = procdim.back();
+            Long w0 = nblocks[light.idim] * heavy.nproc;
+            Long w1 = nblocks[heavy.idim] * light.nproc;
+            if (w0 >= w1) { break; }
+            bool swapped = false;
+            for (auto& f0 : light.procs) {
+                for (auto& f1 : heavy.procs) {
+                    if ((f0 > f1) && (w0*f0 < w1*f1)) {
+                        light.nproc /= f0;
+                        light.nproc *= f1;
+                        heavy.nproc /= f1;
+                        heavy.nproc *= f0;
+                        std::swap(f0,f1);
+                        swapped = true;
+                        break;
+                    }
+                }
+                if (swapped) { break;}
+            }
+            if (!swapped) { break; }
+        }
+
+        for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+            if (!decomp[idim]) {
+                procdim.emplace_back(1,idim);
+            }
+        }
+        for (auto const& pd : procdim) {
+            nprocs[pd.idim] = pd.nproc;
+        }
+    }
+
+    AMREX_ALWAYS_ASSERT(AMREX_D_TERM(nprocs[0],*nprocs[1],*nprocs[2]) == nboxes);
+
+    IntVect const domlo = ccdomain.smallEnd();
+    IntVect const sz = ncells / nprocs;
+    IntVect const extra = ncells - sz*nprocs;
+    auto ixtyp = domain.ixType();
+    BoxList bl(ixtyp);
+#if (AMREX_SPACEDIM == 3)
+    for (int k = 0; k < nprocs[2]; ++k) {
+        // The first extra[2] blocks get one extra cell with a total of
+        // sz[2]+1.  The rest get sz[2] cells.  The decomposition in y
+        // and x directions are similar.
+        int klo = (k < extra[2]) ? k*(sz[2]+1) : (k*sz[2]+extra[2]);
+        int khi = (k < extra[2]) ? klo+(sz[2]+1)-1 : klo+sz[2]-1;
+        klo += domlo[2];
+        khi += domlo[2];
+#endif
+#if (AMREX_SPACEDIM >= 2)
+        for (int j = 0; j < nprocs[1]; ++j) {
+            int jlo = (j < extra[1]) ? j*(sz[1]+1) : (j*sz[1]+extra[1]);
+            int jhi = (j < extra[1]) ? jlo+(sz[1]+1)-1 : jlo+sz[1]-1;
+            jlo += domlo[1];
+            jhi += domlo[1];
+#endif
+            for (int i = 0; i < nprocs[0]; ++i) {
+                int ilo = (i < extra[0]) ? i*(sz[0]+1) : (i*sz[0]+extra[0]);
+                int ihi = (i < extra[0]) ? ilo+(sz[0]+1)-1 : ilo+sz[0]-1;
+                ilo += domlo[0];
+                ihi += domlo[0];
+                Box b(IntVect(AMREX_D_DECL(ilo,jlo,klo)),
+                      IntVect(AMREX_D_DECL(ihi,jhi,khi)));
+                if (b.ok()) {
+                    bl.push_back(b.convert(ixtyp));
+                }
+    AMREX_D_TERM(},},})
+
+    return BoxArray(std::move(bl));
+}
+
 std::ostream&
 operator<< (std::ostream& os, const BoxArray::RefID& id)
 {