From 2c1163b2d244ea8f2d5fa0d8ffa8e3a3141d3b92 Mon Sep 17 00:00:00 2001
From: martinjrobins <martinjrobins@gmail.com>
Date: Mon, 12 Sep 2022 16:04:11 +0100
Subject: [PATCH] #2217 add rest of iterative solvers

---
 .../c_solvers/idaklu/casadi_solver.cpp        | 32 +++++++++++++++++++
 pybamm/solvers/c_solvers/idaklu/common.hpp    |  4 +++
 pybamm/solvers/c_solvers/idaklu/options.cpp   |  5 ++-
 pybamm/solvers/idaklu_solver.py               |  3 +-
 tests/unit/test_solvers/test_idaklu_solver.py |  5 +--
 5 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp
index afacb396ec..919041536a 100644
--- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp
+++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp
@@ -173,6 +173,38 @@ CasadiSolver::CasadiSolver(np_array atol_np, double rel_tol,
     LS = SUNLinSol_SPBCGS(yy, precon_type, options.linsol_max_iterations);
 #endif
   }
+  else if (options.linear_solver == "SUNLinSol_SPFGMR")
+  {
+    DEBUG("\tsetting SUNLinSol_SPFGMR_linear solver");
+#if SUNDIALS_VERSION_MAJOR >= 6
+    LS = SUNLinSol_SPFGMR(yy, precon_type, options.linsol_max_iterations,
+                          sunctx);
+#else
+    LS = SUNLinSol_SPFGMR(yy, precon_type, options.linsol_max_iterations);
+#endif
+  }
+  else if (options.linear_solver == "SUNLinSol_SPGMR")
+  {
+    DEBUG("\tsetting SUNLinSol_SPGMR solver");
+#if SUNDIALS_VERSION_MAJOR >= 6
+    LS = SUNLinSol_SPGMR(yy, precon_type, options.linsol_max_iterations,
+                          sunctx);
+#else
+    LS = SUNLinSol_SPGMR(yy, precon_type, options.linsol_max_iterations);
+#endif
+  }
+  else if (options.linear_solver == "SUNLinSol_SPTFQMR")
+  {
+    DEBUG("\tsetting SUNLinSol_SPGMR solver");
+#if SUNDIALS_VERSION_MAJOR >= 6
+    LS = SUNLinSol_SPTFQMR(yy, precon_type, options.linsol_max_iterations,
+                          sunctx);
+#else
+    LS = SUNLinSol_SPTFQMR(yy, precon_type, options.linsol_max_iterations);
+#endif
+  }
+
+
 
   IDASetLinearSolver(ida_mem, LS, J);
 
diff --git a/pybamm/solvers/c_solvers/idaklu/common.hpp b/pybamm/solvers/c_solvers/idaklu/common.hpp
index 7761778e39..01b4dd4b76 100644
--- a/pybamm/solvers/c_solvers/idaklu/common.hpp
+++ b/pybamm/solvers/c_solvers/idaklu/common.hpp
@@ -18,6 +18,10 @@
 #include <sunlinsol/sunlinsol_dense.h> /* access to dense linear solver          */
 #include <sunlinsol/sunlinsol_lapackdense.h> /* access to lapack linear solver          */
 #include <sunlinsol/sunlinsol_spbcgs.h> /* access to spbcgs iterative linear solver          */
+#include <sunlinsol/sunlinsol_spfgmr.h>
+#include <sunlinsol/sunlinsol_spgmr.h>
+#include <sunlinsol/sunlinsol_sptfqmr.h>
+
 #include <sunmatrix/sunmatrix_sparse.h> /* access to sparse SUNMatrix           */
 #include <sunmatrix/sunmatrix_dense.h> /* access to dense SUNMatrix           */
 
diff --git a/pybamm/solvers/c_solvers/idaklu/options.cpp b/pybamm/solvers/c_solvers/idaklu/options.cpp
index 08ea779c2e..0bafc9ccc7 100644
--- a/pybamm/solvers/c_solvers/idaklu/options.cpp
+++ b/pybamm/solvers/c_solvers/idaklu/options.cpp
@@ -35,7 +35,10 @@ Options::Options(py::dict options)
   else if (linear_solver == "SUNLinSol_KLU" && jacobian == "sparse")
   {
   }
-  else if (linear_solver == "SUNLinSol_SPBCGS" &&
+  else if ((linear_solver == "SUNLinSol_SPBCGS" ||
+            linear_solver == "SUNLinSol_SPFGMR" ||
+            linear_solver == "SUNLinSol_SPGMR" ||
+            linear_solver == "SUNLinSol_SPTFQMR") &&
            (jacobian == "sparse" || jacobian == "matrix-free"))
   {
     using_iterative_solver = true;
diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py
index c6582ec8c9..ad9cc47534 100644
--- a/pybamm/solvers/idaklu_solver.py
+++ b/pybamm/solvers/idaklu_solver.py
@@ -52,7 +52,8 @@ class IDAKLUSolver(pybamm.BaseSolver):
             linear_solver: "SUNLinSol_KLU", # name of sundials linear solver to use
                                   # options are: "SUNLinSol_KLU",
                                   # "SUNLinSol_Dense", "SUNLinSol_LapackDense"
-                                  # "SUNLinSol_SPBCGS"
+                                  # "SUNLinSol_SPBCGS", "SUNLinSol_SPFGMR",
+                                  # "SUNLinSol_SPGMR", "SUNLinSol_SPTFQMR",
             preconditioner: "BBDP", # preconditioner for iterative solvers,
                                     # can be "none", "BBDP"
             linsol_max_iterations: 5, # for iterative linear solvers, max number of
diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py
index 7da3f98ea7..8b36cf394a 100644
--- a/tests/unit/test_solvers/test_idaklu_solver.py
+++ b/tests/unit/test_solvers/test_idaklu_solver.py
@@ -416,7 +416,8 @@ def test_options(self):
         for jacobian in ["none", "dense", "sparse", "matrix-free", "garbage"]:
                 for linear_solver in [
                         "SUNLinSol_SPBCGS", "SUNLinSol_Dense", "SUNLinSol_LapackDense",
-                        "SUNLinSol_KLU", "garbage"
+                        "SUNLinSol_KLU", "SUNLinSol_SPFGMR", "SUNLinSol_SPGMR",
+                        "SUNLinSol_SPTFQMR", "garbage"
                 ]:
                     for precon in ["none", "BBDP"]:
                         options = {
@@ -426,7 +427,7 @@ def test_options(self):
                         }
                         solver = pybamm.IDAKLUSolver(options=options)
                         soln = solver.solve(model, t_eval)
-                        np.testing.assert_array_almost_equal(soln.y, soln_base.y)
+                        np.testing.assert_array_almost_equal(soln.y, soln_base.y, 5)