LinearSolver: fix zero initial guess and update after error

Author: pbrubeck

firedrake/linear_solver.py

@@ -4,6 +4,8 @@
 from firedrake.petsc import PETSc
 from pyop2.mpi import internal_comm
 from firedrake.variational_solver import LinearVariationalProblem, LinearVariationalSolver
+from firedrake.solving_utils import ConvergenceError
+
 
 __all__ = ["LinearSolver"]
 
@@ -83,7 +85,14 @@ def solve(self, x, b):
         if b.function_space() != self.b.function_space():
             raise ValueError(f"b must be a Cofunction in {self.b.function_space()}.")
 
-        self.x.assign(x)
         self.b.assign(b)
-        super().solve()
-        x.assign(self.x)
+        if self.ksp.getInitialGuessNonzero():
+            self.x.assign(x)
+        else:
+            self.x.zero()
+        try:
+            super().solve()
+            x.assign(self.x)
+        except ConvergenceError as e:
+            x.assign(self.x)
+            raise e

tests/firedrake/regression/test_linear_solver.py

@@ -0,0 +1,50 @@
+from firedrake import *
+from firedrake.petsc import PETSc
+import numpy
+
+
+def test_zero_initial_guess():
+    mesh = UnitIntervalMesh(10)
+    space = FunctionSpace(mesh, "CG", 1)
+    test = TestFunction(space)
+    trial = TrialFunction(space)
+
+    solver = LinearSolver(assemble(inner(trial, test) * dx),
+                          solver_parameters={"ksp_type": "preonly",
+                                             "pc_type": "jacobi",
+                                             "ksp_max_it": 1,
+                                             "ksp_initial_guess_nonzero": False})
+    b = assemble(inner(Constant(1), test) * dx)
+
+    u1 = Function(space, name="u1")
+    u1.assign(0)
+    solver.solve(u1, b)
+
+    u2 = Function(space, name="u2")
+    u2.assign(1)
+    solver.solve(u2, b)
+    assert numpy.allclose(u1.dat.data_ro, u2.dat.data_ro)
+
+
+def test_convergence_error_update():
+    mesh = UnitIntervalMesh(10)
+    space = FunctionSpace(mesh, "CG", 1)
+    test = TestFunction(space)
+    trial = TrialFunction(space)
+
+    solver = LinearSolver(assemble(inner(trial, test) * dx),
+                          solver_parameters={"ksp_type": "cg",
+                                             "pc_type": "none",
+                                             "ksp_max_it": 1,
+                                             "ksp_atol": 1.0e-2})
+    b = assemble(inner(Constant(1), test) * dx)
+
+    u = Function(space, name="u")
+    u.assign(-1)
+    uinit = Function(u, name="uinit")
+    try:
+        solver.solve(u, b)
+    except firedrake.exceptions.ConvergenceError:
+        assert solver.ksp.getConvergedReason() == PETSc.KSP.ConvergedReason.DIVERGED_MAX_IT
+
+    assert not numpy.allclose(u.dat.data_ro, uinit.dat.data_ro)