[Breaking] Linear Solver Updates

example/poisson_gmg/plot_convergence.py

@@ -0,0 +1,53 @@
+# =========================================================================================
+# (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
+#
+# This program was produced under U.S. Government contract 89233218CNA000001 for Los
+# Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
+# for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+# in the program are reserved by Triad National Security, LLC, and the U.S. Department
+# of Energy/National Nuclear Security Administration. The Government is granted for
+# itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+# license in this material to reproduce, prepare derivative works, distribute copies to
+# the public, perform publicly and display publicly, and to permit others to do so.
+# =========================================================================================
+
+import numpy as np
+import glob
+import matplotlib.pyplot as plt
+import subprocess
+
+plt.style.use("tableau-colorblind10")
+
+solver = "BiCGSTAB"
+difco = 1e6
+for bound_pro in ["Constant", "Linear"]:
+    for interior_pro in ["Constant", "OldLinear"]:
+
+        p = subprocess.run(
+            [
+                "./poisson-gmg-example",
+                "-i",
+                "parthinput.poisson",
+                "poisson/solver=" + solver,
+                "poisson/interior_D=" + str(difco),
+                "poisson/prolongation=" + bound_pro,
+                "poisson/solver_params/prolongation=" + interior_pro,
+            ],
+            capture_output=True,
+            text=True,
+        )
+        dat = np.genfromtxt(p.stdout.splitlines())
+
+        plt.semilogy(
+            dat[:, 0],
+            dat[:, 1],
+            label=solver + "_" + str(difco) + "_" + bound_pro + "_" + interior_pro,
+        )
+
+
+plt.legend(loc="upper right")
+plt.ylim([1.0e-14, 1e2])
+plt.xlim([0, 40])
+plt.xlabel("# of V-cycles")
+plt.ylabel("RMS Residual")
+plt.savefig("convergence_1e6.pdf")

example/poisson_gmg/poisson_driver.cpp

@@ -75,7 +75,6 @@ TaskCollection PoissonDriver::MakeTaskCollection(BlockList_t &blocks) {
 
   auto pkg = pmesh->packages.Get("poisson_package");
   auto solver = pkg->Param<std::string>("solver");
-  auto flux_correct = pkg->Param<bool>("flux_correct");
   auto use_exact_rhs = pkg->Param<bool>("use_exact_rhs");
   auto *mg_solver =
       pkg->MutableParam<parthenon::solvers::MGSolver<u, rhs, PoissonEquation>>(
@@ -100,9 +99,8 @@ TaskCollection PoissonDriver::MakeTaskCollection(BlockList_t &blocks) {
     if (use_exact_rhs) {
       auto copy_exact = tl.AddTask(get_rhs, TF(solvers::utils::CopyData<exact, u>), md);
       auto comm = AddBoundaryExchangeTasks<BoundaryType::any>(copy_exact, tl, md, true);
-      PoissonEquation eqs;
-      eqs.do_flux_cor = flux_correct;
-      get_rhs = eqs.Ax<u, rhs>(tl, comm, md);
+      auto *eqs = pkg->MutableParam<PoissonEquation>("poisson_equation");
+      get_rhs = eqs->Ax<u, rhs>(tl, comm, md);
     }
 
     // Set initial solution guess to zero

example/poisson_gmg/poisson_equation.hpp

@@ -36,6 +36,26 @@ class PoissonEquation {
  public:
   bool do_flux_cor = false;
   bool set_flux_boundary = false;
+  bool include_flux_dx = false;
+  enum class ProlongationType { Constant, Linear, Kwak };
+  ProlongationType prolongation_type = ProlongationType::Constant;
+
+  PoissonEquation(parthenon::ParameterInput *pin, const std::string &label) {
+    do_flux_cor = pin->GetOrAddBoolean(label, "flux_correct", false);
+    set_flux_boundary = pin->GetOrAddBoolean(label, "set_flux_boundary", false);
+    include_flux_dx =
+        (pin->GetOrAddString(label, "boundary_prolongation", "Linear") == "Constant");
+    auto pro_int = pin->GetOrAddString(label, "interior_prolongation", "Linear");
+    if (pro_int == "Constant") {
+      prolongation_type = ProlongationType::Constant;
+    } else if (pro_int == "Linear") {
+      prolongation_type = ProlongationType::Linear;
+    } else if (pro_int == "Kwak") {
+      prolongation_type = ProlongationType::Kwak;
+    } else {
+      PARTHENON_FAIL("Invalid user prolongation type.");
+    }
+  }
 
   // Add tasks to calculate the result of the matrix A (which is implicitly defined by
   // this class) being applied to x_t and store it in field out_t
@@ -44,7 +64,7 @@ class PoissonEquation {
                        std::shared_ptr<parthenon::MeshData<Real>> &md) {
     auto flux_res = tl.AddTask(depends_on, CalculateFluxes<x_t>, md);
     if (set_flux_boundary) {
-      flux_res = tl.AddTask(flux_res, SetFluxBoundaries<x_t>, md);
+      flux_res = tl.AddTask(flux_res, SetFluxBoundaries<x_t>, md, include_flux_dx);
     }
     if (do_flux_cor && !(md->grid.type == parthenon::GridType::two_level_composite)) {
       auto start_flxcor =
@@ -164,9 +184,128 @@ class PoissonEquation {
     return TaskStatus::complete;
   }
 
+  template <class... var_ts>
+  parthenon::TaskID Prolongate(parthenon::TaskList &tl, parthenon::TaskID depends_on,
+                               std::shared_ptr<parthenon::MeshData<Real>> &md) {
+    if (prolongation_type == ProlongationType::Constant) {
+      return tl.AddTask(depends_on, ProlongateImpl<ProlongationType::Constant, var_ts...>,
+                        md);
+    } else if (prolongation_type == ProlongationType::Linear) {
+      return tl.AddTask(depends_on, ProlongateImpl<ProlongationType::Linear, var_ts...>,
+                        md);
+    } else if (prolongation_type == ProlongationType::Kwak) {
+      return tl.AddTask(depends_on, ProlongateImpl<ProlongationType::Kwak, var_ts...>,
+                        md);
+    }
+    return depends_on;
+  }
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  static Real LinearFactor(int d, bool lo_bound, bool up_bound) {
+    if (d == 0) return 1.0; // Indicates this dimension is not included
+    if (d == 1) return (2.0 + !up_bound) / 4.0;
+    if (d == -1) return (2.0 + !lo_bound) / 4.0;
+    if (d == 3) return !up_bound / 4.0;
+    if (d == -3) return !lo_bound / 4.0;
+    return 0.0;
+  }
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  static Real QuadraticFactor(int d) {
+    if (d == 0) return 1.0; // Indicates this dimension is not included
+    if (d == 1 || d == -1) return 30.0 / 32.0;
+    if (d == 3 || d == -3) return 5.0 / 32.0;
+    if (d == 5 || d == -5) return -3.0 / 32.0;
+    return 0.0;
+  }
+
+  template <ProlongationType prolongation_type, class... var_ts>
+  static parthenon::TaskStatus
+  ProlongateImpl(std::shared_ptr<parthenon::MeshData<Real>> &md) {
+    using namespace parthenon;
+    const int ndim = md->GetMeshPointer()->ndim;
+    IndexRange ib = md->GetBoundsI(IndexDomain::interior);
+    IndexRange jb = md->GetBoundsJ(IndexDomain::interior);
+    IndexRange kb = md->GetBoundsK(IndexDomain::interior);
+    IndexRange cib = md->GetBoundsI(CellLevel::coarse, IndexDomain::interior);
+    IndexRange cjb = md->GetBoundsJ(CellLevel::coarse, IndexDomain::interior);
+    IndexRange ckb = md->GetBoundsK(CellLevel::coarse, IndexDomain::interior);
+
+    using TE = parthenon::TopologicalElement;
+
+    int nblocks = md->NumBlocks();
+    std::vector<bool> include_block(nblocks, true);
+    for (int b = 0; b < nblocks; ++b) {
+      include_block[b] =
+          md->grid.logical_level == md->GetBlockData(b)->GetBlockPointer()->loc.level();
+    }
+    const auto desc = parthenon::MakePackDescriptor<var_ts...>(md.get());
+    const auto desc_coarse =
+        parthenon::MakePackDescriptor<var_ts...>(md.get(), {}, {PDOpt::Coarse});
+    auto pack = desc.GetPack(md.get(), include_block);
+    auto pack_coarse = desc_coarse.GetPack(md.get(), include_block);
+
+    parthenon::par_for(
+        "Prolongate", 0, pack.GetNBlocks() - 1, pack.GetLowerBoundHost(0),
+        pack.GetUpperBoundHost(0), kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+        KOKKOS_LAMBDA(const int b, const int n, const int fk, const int fj,
+                      const int fi) {
+          const int ck = (ndim > 2) ? (fk - kb.s) / 2 + ckb.s : ckb.s;
+          const int cj = (ndim > 1) ? (fj - jb.s) / 2 + cjb.s : cjb.s;
+          const int ci = (ndim > 0) ? (fi - ib.s) / 2 + cib.s : cib.s;
+          const int fok = (fk - kb.s) % 2;
+          const int foj = (fj - jb.s) % 2;
+          const int foi = (fi - ib.s) % 2;
+          const bool bound[6]{pack.IsPhysicalBoundary(b, 0, 0, -1) && (ib.s == fi),
+                              pack.IsPhysicalBoundary(b, 0, 0, 1) && (ib.e == fi),
+                              pack.IsPhysicalBoundary(b, 0, -1, 0) && (jb.s == fj),
+                              pack.IsPhysicalBoundary(b, 0, 1, 0) && (jb.e == fj),
+                              pack.IsPhysicalBoundary(b, -1, 0, 0) && (kb.s == fk),
+                              pack.IsPhysicalBoundary(b, 1, 0, 0) && (kb.e == fk)};
+
+          if constexpr (ProlongationType::Constant == prolongation_type) {
+            pack(b, n, fk, fj, fi) = pack_coarse(b, n, ck, cj, ci);
+          } else if constexpr (ProlongationType::Linear == prolongation_type) {
+            pack(b, n, fk, fj, fi) = 0.0;
+            for (int ok = -(ndim > 2); ok < 1 + (ndim > 2); ++ok) {
+              for (int oj = -(ndim > 1); oj < 1 + (ndim > 1); ++oj) {
+                for (int oi = -(ndim > 0); oi < 1 + (ndim > 0); ++oi) {
+                  const int dx3 = (ndim > 2) ? 4 * ok - (2 * fok - 1) : 0;
+                  const int dx2 = (ndim > 1) ? 4 * oj - (2 * foj - 1) : 0;
+                  const int dx1 = 4 * oi - (2 * foi - 1);
+                  pack(b, n, fk, fj, fi) += LinearFactor(dx1, bound[0], bound[1]) *
+                                            LinearFactor(dx2, bound[2], bound[3]) *
+                                            LinearFactor(dx3, bound[4], bound[5]) *
+                                            pack_coarse(b, n, ck + ok, cj + oj, ci + oi);
+                }
+              }
+            }
+          } else if constexpr (ProlongationType::Kwak == prolongation_type) {
+            pack(b, n, fk, fj, fi) = 0.0;
+            if (ndim > 2 && !bound[4 + fok]) {
+              for (int ok = fok - 1; ok <= fok; ++ok) {
+                pack(b, n, fk, fj, fi) += pack_coarse(b, n, ck + ok, cj, ci);
+              }
+            }
+            if (ndim > 1 && !bound[2 + foj]) {
+              for (int oj = foj - 1; oj <= foj; ++oj) {
+                pack(b, n, fk, fj, fi) += pack_coarse(b, n, ck, cj + oj, ci);
+              }
+            }
+            if (ndim > 0 && !bound[foi]) {
+              for (int oi = foi - 1; oi <= foi; ++oi) {
+                pack(b, n, fk, fj, fi) += pack_coarse(b, n, ck, cj, ci + oi);
+              }
+            }
+            pack(b, n, fk, fj, fi) /= 2.0 * ndim;
+          }
+        });
+    return TaskStatus::complete;
+  }
+
   template <class var_t>
   static parthenon::TaskStatus
-  SetFluxBoundaries(std::shared_ptr<parthenon::MeshData<Real>> &md) {
+  SetFluxBoundaries(std::shared_ptr<parthenon::MeshData<Real>> &md, bool do_flux_dx) {
     using namespace parthenon;
     const int ndim = md->GetMeshPointer()->ndim;
     IndexRange ib = md->GetBoundsI(IndexDomain::interior);
@@ -196,7 +335,6 @@ class PoissonEquation {
     constexpr int x3off[6]{0, 0, 0, 0, -1, 1};
     constexpr TE tes[6]{TE::F1, TE::F1, TE::F2, TE::F2, TE::F3, TE::F3};
     constexpr int dirs[6]{X1DIR, X1DIR, X2DIR, X2DIR, X3DIR, X3DIR};
-
     parthenon::par_for_outer(
         DEFAULT_OUTER_LOOP_PATTERN, "SetFluxBoundaries", DevExecSpace(),
         scratch_size_in_bytes, scratch_level, 0, pack.GetNBlocks() - 1,
@@ -228,7 +366,8 @@ class PoissonEquation {
             }
             // Correct for size of neighboring zone at fine-coarse boundary when using
             // constant prolongation
-            if (pack.GetLevel(b, x3off[face], x2off[face], x1off[face]) == level - 1) {
+            if (do_flux_dx &&
+                pack.GetLevel(b, x3off[face], x2off[face], x1off[face]) == level - 1) {
               parthenon::par_for_inner(DEFAULT_INNER_LOOP_PATTERN, member, 0,
                                        idxer.size() - 1, [&](const int idx) {
                                          const auto [k, j, i] = idxer(idx);

example/poisson_gmg/poisson_package.cpp

@@ -78,24 +78,19 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   pkg->UserBoundaryFunctions[BF::outer_x2].push_back(GetBC<X2DIR, BCSide::Outer>());
   pkg->UserBoundaryFunctions[BF::outer_x3].push_back(GetBC<X3DIR, BCSide::Outer>());
 
-  int max_poisson_iterations = pin->GetOrAddInteger("poisson", "max_iterations", 10000);
-  pkg->AddParam<>("max_iterations", max_poisson_iterations);
-
   Real diagonal_alpha = pin->GetOrAddReal("poisson", "diagonal_alpha", 0.0);
   pkg->AddParam<>("diagonal_alpha", diagonal_alpha);
 
   std::string solver = pin->GetOrAddString("poisson", "solver", "MG");
   pkg->AddParam<>("solver", solver);
 
-  Real err_tol = pin->GetOrAddReal("poisson", "error_tolerance", 1.e-8);
-  pkg->AddParam<>("error_tolerance", err_tol);
-
   bool use_exact_rhs = pin->GetOrAddBoolean("poisson", "use_exact_rhs", false);
   pkg->AddParam<>("use_exact_rhs", use_exact_rhs);
 
-  PoissonEquation eq;
-  eq.do_flux_cor = pin->GetOrAddBoolean("poisson", "flux_correct", false);
-  eq.set_flux_boundary = pin->GetOrAddBoolean("poisson", "set_flux_boundary", false);
+  std::string prolong = pin->GetOrAddString("poisson", "boundary_prolongation", "Linear");
+
+  PoissonEquation eq(pin, "poisson");
+  pkg->AddParam<>("poisson_equation", eq, parthenon::Params::Mutability::Mutable);
 
   parthenon::solvers::MGParams mg_params(pin, "poisson/solver_params");
   parthenon::solvers::MGSolver<u, rhs, PoissonEquation> mg_solver(pkg.get(), mg_params,
@@ -124,7 +119,6 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 
   auto mflux_comm = Metadata({Metadata::Cell, Metadata::Independent, Metadata::FillGhost,
                               Metadata::WithFluxes, Metadata::GMGRestrict});
-  std::string prolong = pin->GetOrAddString("poisson", "prolongation", "Linear");
   if (prolong == "Linear") {
     mflux_comm.RegisterRefinementOps<ProlongateSharedLinear, RestrictAverage>();
   } else if (prolong == "Constant") {

src/prolong_restrict/pr_ops.hpp

@@ -288,104 +288,6 @@ using ProlongateSharedMinMod = ProlongateSharedGeneral<true, false>;
 using ProlongateSharedLinear = ProlongateSharedGeneral<false, false>;
 using ProlongatePiecewiseConstant = ProlongateSharedGeneral<false, true>;
 
-enum class MGProlongationType { Constant, Linear, Quadratic, Kwak };
-
-template <MGProlongationType type>
-struct ProlongateSharedMG {
-  static constexpr bool OperationRequired(TopologicalElement fel,
-                                          TopologicalElement cel) {
-    if (fel != TopologicalElement::CC) return false;
-    return fel == cel;
-  }
-
-  KOKKOS_FORCEINLINE_FUNCTION
-  static Real QuadraticFactor(int d) {
-    if (d == 0) return 1.0; // Indicates this dimension is not included
-    if (d == 1 || d == -1) return 30.0 / 32.0;
-    if (d == 3 || d == -3) return 5.0 / 32.0;
-    if (d == 5 || d == -5) return -3.0 / 32.0;
-    return 0.0;
-  }
-
-  KOKKOS_FORCEINLINE_FUNCTION
-  static Real LinearFactor(int d, bool up_bound, bool lo_bound) {
-    if (d == 0) return 1.0; // Indicates this dimension is not included
-    if (d == 1) return (2.0 + !up_bound) / 4.0;
-    if (d == -1) return (2.0 + !lo_bound) / 4.0;
-    if (d == 3) return !up_bound / 4.0;
-    if (d == -3) return !lo_bound / 4.0;
-    return 0.0;
-  }
-
-  KOKKOS_FORCEINLINE_FUNCTION
-  static Real ConstantFactor(int d) {
-    if (d == 0) return 1.0; // Indicates this dimension is not included
-    if (d == 1 || d == -1) return 1.0;
-    return 0.0;
-  }
-
-  template <int DIM, TopologicalElement el = TopologicalElement::CC,
-            TopologicalElement /*cel*/ = TopologicalElement::CC>
-  KOKKOS_FORCEINLINE_FUNCTION static void
-  Do(const int l, const int m, const int n, const int k, const int j, const int i,
-     const IndexRange &ckb, const IndexRange &cjb, const IndexRange &cib,
-     const IndexRange &kb, const IndexRange &jb, const IndexRange &ib,
-     const Coordinates_t &coords, const Coordinates_t &coarse_coords,
-     const ParArrayND<Real, VariableState> *pcoarse,
-     const ParArrayND<Real, VariableState> *pfine) {
-    using namespace util;
-    auto &coarse = *pcoarse;
-    auto &fine = *pfine;
-
-    constexpr int element_idx = static_cast<int>(el) % 3;
-
-    const int fi = (DIM > 0) ? (i - cib.s) * 2 + ib.s : ib.s;
-    const int fj = (DIM > 1) ? (j - cjb.s) * 2 + jb.s : jb.s;
-    const int fk = (DIM > 2) ? (k - ckb.s) * 2 + kb.s : kb.s;
-
-    for (int fok = 0; fok < 1 + (DIM > 2); ++fok) {
-      for (int foj = 0; foj < 1 + (DIM > 1); ++foj) {
-        for (int foi = 0; foi < 1 + (DIM > 0); ++foi) {
-          auto &f = fine(element_idx, l, m, n, fk + fok, fj + foj, fi + foi);
-          f = 0.0;
-          const bool lo_bound_x = ((fi + foi) == ib.s);
-          const bool up_bound_x = ((fi + foi) == ib.e);
-          const bool lo_bound_y = ((fj + foj) == jb.s);
-          const bool up_bound_y = ((fj + foj) == jb.e);
-          const bool lo_bound_z = ((fk + fok) == kb.s);
-          const bool up_bound_z = ((fk + fok) == kb.e);
-          for (int ok = -(DIM > 2); ok < 1 + (DIM > 2); ++ok) {
-            for (int oj = -(DIM > 1); oj < 1 + (DIM > 1); ++oj) {
-              for (int oi = -(DIM > 0); oi < 1 + (DIM > 0); ++oi) {
-                const int dx = 4 * oi - foi + 1;
-                const int dy = (DIM > 1) ? 4 * oj - foj + 1 : 0;
-                const int dz = (DIM > 2) ? 4 * ok - fok + 1 : 0;
-                if constexpr (MGProlongationType::Linear == type) {
-                  f += LinearFactor(dx, lo_bound_x, up_bound_x) *
-                       LinearFactor(dy, lo_bound_y, up_bound_y) *
-                       LinearFactor(dz, lo_bound_z, up_bound_z) *
-                       coarse(element_idx, l, m, n, k + ok, j + oj, i + oi);
-                } else if constexpr (MGProlongationType::Kwak == type) {
-                  const Real fac =
-                      ((dx <= 1) + (dy <= 1 && DIM > 1) + (dz <= 1 && DIM > 2)) /
-                      (2.0 * DIM);
-                  f += fac * coarse(element_idx, l, m, n, k + ok, j + oj, i + oi);
-                } else if constexpr (MGProlongationType::Quadratic == type) {
-                  f += QuadraticFactor(dx) * QuadraticFactor(dy) * QuadraticFactor(dz) *
-                       coarse(element_idx, l, m, n, k + ok, j + oj, i + oi);
-                } else {
-                  f += ConstantFactor(dx) * ConstantFactor(dy) * ConstantFactor(dz) *
-                       coarse(element_idx, l, m, n, k + ok, j + oj, i + oi);
-                }
-              }
-            }
-          }
-        }
-      }
-    }
-  }
-};
-
 struct ProlongateInternalAverage {
   static constexpr bool OperationRequired(TopologicalElement fel,
                                           TopologicalElement cel) {