[Breaking] Linear Solver Updates

CHANGELOG.md

@@ -36,6 +36,7 @@
 
 
 ### Incompatibilities (i.e. breaking changes)
+- [[PR 1174]](https://github.com/parthenon-hpc-lab/parthenon/pull/1174) Add CG solver, custom solver prolongation operator options, and switch to stage based solvers
 - [[PR 1177]](https://github.com/parthenon-hpc-lab/parthenon/pull/1177) Make mesh-level boundary conditions usable without the "user" flag
 
 ## Release 24.08

doc/sphinx/src/solvers.rst

@@ -3,23 +3,85 @@
 Solvers
 =======
 
-Parthenon does not yet provide an exhaustive set of plug and play solvers. 
-Nevertheless, the building blocks required for implementing Krylov subspace 
-methods (i.e. global reductions for vector dot products) like CG, BiCGStab, 
-and GMRES are available. An example of a Parthenon based implementation of 
-BiCGStab can be found in ``examples/poisson_gmg``. Additionally, the 
-infrastructure required for implementing multigrid solvers is also 
-included in Parthenon. The requisite hierarchy of grids is produced if 
-``parthenon/mesh/multigrid=true`` is set in the parameter input. An example 
-of a multi-grid based linear solver in Parthenon is also given in 
-``examples/poisson_gmg`` (and also an example of using multi-grid as a 
-preconditioner for BiCGStab). We plan to build wrappers that simplify the 
-use of these methods in down stream codes in the future. Note that the 
-example code does not currently rely on the Stencil and SparseMatrixAccessor 
-code described below.
+Parthenon provides a number of linear solvers, including a geometric multigrid
+solver, a CG solver, a BiCGSTAB solver, and multigrid preconditioned versions
+of the latter two solvers. 
+
+Solvers are templated on a type defining the system of equations they are solving. 
+The type defining the system of equations must provide two methods and a ``TypeList`` 
+of all of the fields that make up the vector space:
+.. code:: c++
+  class MySystemOfEquations {
+    using IndependentVars = parthenon::TypeList<var1_t, var2_t>;
+
+    TaskId Ax(TaskList &tl, TaskID depends_on,
+              std::shared_ptr<parthenon::MeshData<Real>> &md_mat,
+              std::shared_ptr<parthenon::MeshData<Real>> &md_in,
+              std::shared_ptr<parthenon::MeshData<Real>> &md_out);
+
+    TaskStatus SetDiagonal(std::shared_ptr<parthenon::MeshData<Real>> &md_mat,
+                           std::shared_ptr<parthenon::MeshData<Real>> &md_diag)
+  };
+
+The routine ``Ax`` must calculate the matrix vector product ``y <- A.x`` by taking a container
+``md_mat`` which contains all of the fields required to reconstruct the matrix ``A`` associated
+with the system of linear equations, the container ``md_in`` which will store the vector ``x`` 
+in the fields in the typelist ``IndependentVars``, and ``md_out`` which will hold the vector ``y``. 
+
+The routine ``SetDiagonal`` takes the same container ``md_mat`` as ``Ax`` and returns the
+(approximate) diagonal of ``A`` in the container ``md_diag``. This only needs to be approximate
+since it is only used in preconditioners/smoothers.
+
+With such a class defining a linear system of equations, one can then define and use a solver with 
+code along the lines of:
+.. code:: c++ 
+  std::string base_cont_name = "base";
+  std::string u_cont_name = "u";
+  std::string rhs_cont_name = "rhs";
+
+  MySystemOfEquations eqs(....);
+  std::shared_ptr<SolverBase> psolver = std::make_shared<BiCGSTABSolver<MySystemOfEquations>>(
+      base_cont_name, u_cont_name, rhs_cont_name, pin, "location/of/solver_params", eqs);
+
+  ...
+
+  auto partitions = pmesh->GetDefaultBlockPartitions();
+  const int num_partitions = partitions.size();
+  TaskRegion &region = tc.AddRegion(num_partitions);
+  for (int partition = 0; partition < num_partitions; ++partition) {
+    TaskList &tl = region[partition];
+    auto &md = pmesh->mesh_data.Add(base_cont_name, partitions[partition]);
+    auto &md_u = pmesh->mesh_data.Add(u_cont_name, md);
+    auto &md_rhs = pmesh->mesh_data.Add(rhs_cont_name, md);
+
+    // Do some stuff to fill the base container with information necessary to define A
+    // if it wasn't defined during initialization or something
+
+    // Do some stuff to fill the rhs container
+
+    auto setup = psolver->AddSetupTasks(tl, dependence, partition, pmesh);
+    auto solve = psolver->AddTasks(tl, setup, partition, pmesh); 
+
+    // Do some stuff with the solution stored in md_u
+  }
+
+Some notes: 
+- All solvers inherit from ``SolverBase``, so the best practice is to stash a shared pointer to a 
+  ``SolverBase`` object in params during initialization and pull this solver out while building a 
+  task list. This should make switching between solvers trivial.
+- For any solver involving geometric multigrid, the input parameter 
+  ``parthenon/mesh/multigrid`` must be set to ``true``. This tells the ``Mesh``
+  to build the coarsened blocks associated with the multi-grid hierarchy.
+- For geometric multigrid based solvers, it is possible to define block interior prolongation 
+  operators that are separate from the standard prolongation machinery in Parthenon. This allows 
+  for defining boundary aware prolongation operators and having different prolongation operators
+  in the ghost cells of blocks from the prolongation operators used in their interiors. Users can 
+  easily define their own prolongation operators. The prolongation functor is passed as a template 
+  argument to the multi-grid solver class. An example of using these interior prolongation
+  operators is contained in the ``poisson_gmg`` example.
 
 Some implementation notes about geometric multi-grid can be found in 
-:ref:`these notes <doc/latex/main.pdf>`. 
+:ref:`these notes <doc/latex/main.pdf>`.
 
 Stencil
 -------

example/poisson_gmg/plot_convergence.py

@@ -0,0 +1,66 @@
+# =========================================================================================
+# (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"
+solver_lbl = "BCGS"
+difco = 1e6
+refine = True
+
+for bound_pro in ["Constant", "Linear"]:
+    for interior_pro in ["Constant", "Linear", "SplitLin", "Kwak"]:
+        command = ["./poisson-gmg-example", "-i", "parthinput.poisson"]
+        command.append("poisson/solver=" + solver)
+        command.append("poisson/interior_D=" + str(difco))
+        command.append("poisson/boundary_prolongation=" + bound_pro)
+        if interior_pro == "SplitLin":
+            command.append("poisson/solver_params/prolongation=Default")
+        else:
+            command.append("poisson/interior_prolongation=" + interior_pro)
+            command.append("poisson/solver_params/prolongation=User")
+
+        if refine:
+            command.append("parthenon/static_refinement0/x1min=-1.0")
+            command.append("parthenon/static_refinement0/x1max=-0.75")
+            command.append("parthenon/static_refinement0/x2min=-1.0")
+            command.append("parthenon/static_refinement0/x2max=-0.75")
+            command.append("parthenon/static_refinement0/level=3")
+
+        p = subprocess.run(command, capture_output=True, text=True)
+        lines = p.stdout.splitlines()
+        # Ignore any initialization junk that gets spit out earlier from adding parameters
+        idx = lines.index("# [0] v-cycle")
+        dat = np.genfromtxt(lines[idx:])
+        label = "{}_{}".format(solver_lbl, interior_pro)
+        if refine:
+            label = "{}_{}_Bnd{}".format(solver_lbl, interior_pro, bound_pro)
+        plt.semilogy(
+            dat[:, 0],
+            dat[:, 1],
+            label=label.replace("Constant", "Const").replace("Linear", "Lin"),
+        )
+
+
+plt.legend()
+plt.ylim([1.0e-14, 1e4])
+plt.xlim([0, 40])
+plt.xlabel("# of V-cycles")
+plt.ylabel("RMS Residual")
+plt.title("$D_{int} = 10^6$ w/ Refinement")
+plt.savefig("convergence_1e6.pdf")

example/poisson_gmg/poisson_driver.cpp

@@ -29,7 +29,9 @@
 #include "poisson_package.hpp"
 #include "prolong_restrict/prolong_restrict.hpp"
 #include "solvers/bicgstab_solver.hpp"
+#include "solvers/cg_solver.hpp"
 #include "solvers/mg_solver.hpp"
+#include "solvers/solver_utils.hpp"
 
 using namespace parthenon::driver::prelude;
 
@@ -45,18 +47,9 @@ parthenon::DriverStatus PoissonDriver::Execute() {
 
   // After running, retrieve the final residual for checking in tests
   auto pkg = pmesh->packages.Get("poisson_package");
-  auto solver = pkg->Param<std::string>("solver");
-  if (solver == "BiCGSTAB") {
-    auto *bicgstab_solver =
-        pkg->MutableParam<parthenon::solvers::BiCGSTABSolver<u, rhs, PoissonEquation>>(
-            "MGBiCGSTABsolver");
-    final_rms_residual = bicgstab_solver->GetFinalResidual();
-  } else if (solver == "MG") {
-    auto *mg_solver =
-        pkg->MutableParam<parthenon::solvers::MGSolver<u, rhs, PoissonEquation>>(
-            "MGsolver");
-    final_rms_residual = mg_solver->GetFinalResidual();
-  }
+  auto psolver =
+      pkg->Param<std::shared_ptr<parthenon::solvers::SolverBase>>("solver_pointer");
+  final_rms_residual = psolver->GetFinalResidual();
 
   return DriverStatus::complete;
 }
@@ -68,54 +61,53 @@ TaskCollection PoissonDriver::MakeTaskCollection(BlockList_t &blocks) {
   TaskID none(0);
 
   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>>(
-          "MGsolver");
-  auto *bicgstab_solver =
-      pkg->MutableParam<parthenon::solvers::BiCGSTABSolver<u, rhs, PoissonEquation>>(
-          "MGBiCGSTABsolver");
+  auto psolver =
+      pkg->Param<std::shared_ptr<parthenon::solvers::SolverBase>>("solver_pointer");
 
   auto partitions = pmesh->GetDefaultBlockPartitions();
   const int num_partitions = partitions.size();
   TaskRegion &region = tc.AddRegion(num_partitions);
   for (int i = 0; i < num_partitions; ++i) {
     TaskList &tl = region[i];
     auto &md = pmesh->mesh_data.Add("base", partitions[i]);
+    auto &md_u = pmesh->mesh_data.Add("u", md, {u::name()});
+    auto &md_rhs = pmesh->mesh_data.Add("rhs", md, {u::name()});
+
+    // Move the rhs variable into the rhs stage for stage based solver
+    auto copy_rhs =
+        tl.AddTask(none, TF(solvers::utils::between_fields::CopyData<rhs, u>), md);
+    copy_rhs = tl.AddTask(copy_rhs, TF(solvers::utils::CopyData<parthenon::TypeList<u>>),
+                          md, md_rhs);
 
     // Possibly set rhs <- A.u_exact for a given u_exact so that the exact solution is
     // known when we solve A.u = rhs
-    auto get_rhs = none;
     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 copy_exact = tl.AddTask(
+          copy_rhs, TF(solvers::utils::between_fields::CopyData<exact, u>), md);
+      auto copy_u_between_stages = tl.AddTask(
+          copy_exact, TF(solvers::utils::CopyData<parthenon::TypeList<u>>), md, md_u);
+      auto comm = AddBoundaryExchangeTasks<BoundaryType::any>(copy_u_between_stages, tl,
+                                                              md_u, true);
+      auto *eqs =
+          pkg->MutableParam<poisson_package::PoissonEquation<u, D>>("poisson_equation");
+      copy_rhs = eqs->Ax(tl, comm, md, md_u, md_rhs);
     }
 
     // Set initial solution guess to zero
-    auto zero_u = tl.AddTask(get_rhs, TF(solvers::utils::SetToZero<u>), md);
-
-    auto solve = zero_u;
-    if (solver == "BiCGSTAB") {
-      auto setup = bicgstab_solver->AddSetupTasks(tl, zero_u, i, pmesh);
-      solve = bicgstab_solver->AddTasks(tl, setup, pmesh, i);
-    } else if (solver == "MG") {
-      auto setup = mg_solver->AddSetupTasks(tl, zero_u, i, pmesh);
-      solve = mg_solver->AddTasks(tl, setup, pmesh, i);
-    } else {
-      PARTHENON_FAIL("Unknown solver type.");
-    }
+    auto zero_u = tl.AddTask(copy_rhs, TF(solvers::utils::SetToZero<u>), md_u);
+    auto setup = psolver->AddSetupTasks(tl, zero_u, i, pmesh);
+    auto solve = psolver->AddTasks(tl, setup, i, pmesh);
 
     // If we are using a rhs to which we know the exact solution, compare our computed
     // solution to the exact solution
     if (use_exact_rhs) {
-      auto diff = tl.AddTask(solve, TF(solvers::utils::AddFieldsAndStore<exact, u, u>),
-                             md, 1.0, -1.0);
-      auto get_err = solvers::utils::DotProduct<u, u>(diff, tl, &err, md);
+      auto copy_back = tl.AddTask(
+          solve, TF(solvers::utils::CopyData<parthenon::TypeList<u>>), md_u, md);
+      auto diff = tl.AddTask(
+          copy_back, TF(solvers::utils::between_fields::AddFieldsAndStore<exact, u, u>),
+          md, 1.0, -1.0);
+      auto get_err = solvers::utils::between_fields::DotProduct<u, u>(diff, tl, &err, md);
       tl.AddTask(
           get_err,
           [](PoissonDriver *driver, int partition) {