rewriting of the custom PETScLUSolver with a name change to compose w…

…ith the dolfin version and added tests

soupy/solver/PETScLUSolver.py

@@ -0,0 +1,66 @@
+import dolfin as dl 
+from petsc4py import PETSc
+from mpi4py import MPI
+
+
+class PETScLUSolver:
+    """
+    Wrapper for `dolfin.PETScLUSolver` with methods
+    :code:`solve_transpose` and :code:`set_operator`
+    """
+    def __init__(self, mpi_comm=MPI.COMM_WORLD, method="mumps"):
+        self.mpi_comm = mpi_comm 
+        self.method = method 
+        self.solver = dl.PETScLUSolver(mpi_comm, method)
+        self.ksp = self.solver.ksp()
+
+    def set_operator(self, A):
+        if hasattr(A, 'mat'):
+            self.ksp.setOperators(A.mat())
+        else:
+            self.ksp.setOperators(dl.as_backend_type(A).mat())
+
+    def solve(self, x, b):
+        """
+        Solve the linear system 
+
+        .. math:: `Ax = b`
+
+        and stores result to :code:`x`
+        """
+
+        self.solver.solve(x, b)
+
+    def solve_transpose(self, x, b):
+        """
+        Solve the linear system 
+
+        .. math:: `A^T x = b`
+
+        and stores result to :code:`x`
+        """
+        if hasattr(self.solver, 'solve_transpose'):
+            self.solver.solveTranspose(x, b)
+        else:
+            self._custom_solveTranspose(x, b)
+
+
+    def _custom_solveTranspose(self, x, b):
+        if hasattr(b, 'vec'):
+            b_vec = b.vec()
+        else:
+            b_vec = dl.as_backend_type(b).vec()
+
+        if hasattr(x, 'vec'):
+            # x already has vec. Can directly solve to x.vec()
+            self.ksp.solveTranspose(b_vec, x.vec())
+
+        else:
+            # x has no vec object. Need to make one, and then 
+            # store the result to x
+            x_petsc = dl.as_backend_type(x)
+            self.ksp.solveTranspose(b_vec, x_petsc.vec())
+
+            # Now store to the input x
+            x.set_local(x_petsc.get_local())
+            x.apply("")

soupy/solver/__init__.py

@@ -15,4 +15,4 @@
 
 from .nonlinearVariationalSolver import NonlinearVariationalSolver
 
-from .customPETScLUSolver import CustomPETScLUSolver
+from .PETScLUSolver import PETScLUSolver

soupy/test/test_customPETScLUSolver.py

@@ -4,38 +4,57 @@
 import os 
 
 sys.path.append(os.environ.get('HIPPYLIB_PATH'))
+import hippylib as hp 
+
+sys.path.append('../../')
+from soupy import PETScLUSolver
+
+
+def make_linear_system(V, A_tensor, b_tensor):
+    """
+    Assembles the linear system for a linear ADR problem. \
+        This is not symmetric 
+    """
+
+    u_trial = dl.TrialFunction(V)
+    u_test = dl.TestFunction(V)
+
+    a_form = dl.inner(dl.grad(u_trial), dl.grad(u_test))*dl.dx \
+            + dl.inner(dl.Constant((1.0, 1.0)), dl.grad(u_trial))*u_test*dl.dx \
+            + u_trial*u_test*dl.dx
+    b_form = dl.Constant(1.0) * u_test * dl.dx 
+
+    boundary_term = dl.Expression("x[0]", degree=1)
+    bcs = dl.DirichletBC(V, boundary_term, "on_boundary")
+
+    dl.assemble_system(a_form, b_form, bcs=bcs, A_tensor=A_tensor, b_tensor=b_tensor)
 
-sys.path.append("../solver")
-from customPETScLUSolver import CustomPETScLUSolver
 
 class TestCustomPETScLUSolver(unittest.TestCase):
     def setUp(self):
         self.mesh = dl.UnitSquareMesh(20,20)
         self.method = "mumps"
+        self.V = dl.FunctionSpace(self.mesh, "CG", 2)
         self.tol = 1e-12
 
-    def testSolver(self):
-        V = dl.FunctionSpace(self.mesh, "CG", 2)
-        u_trial = dl.TrialFunction(V)
-        u_test = dl.TestFunction(V)
-        u_fun = dl.Function(V)
-
-
-        a_form = dl.inner(dl.grad(u_trial), dl.grad(u_test))*dl.dx + u_trial*u_test*dl.dx
-        b_form = dl.Constant(1.0) * u_test * dl.dx 
-
-        bcs = dl.DirichletBC(V, dl.Constant(0.0), "on_boundary")
 
+    def compareSolve(self, backend="dolfin"):
+        """
+        Compares the solution of using custom solver against \
+            Standard solver
+        """
         A = dl.PETScMatrix()
         b = dl.PETScVector()
-        dl.assemble_system(a_form, b_form, bcs=bcs, A_tensor=A, b_tensor=b)
+
+        make_linear_system(self.V, A, b)
 
+        u_fun = dl.Function(self.V)
         u = u_fun.vector()
         u_custom = u.copy()
 
-        diff = dl.Function(V).vector()
+        diff = dl.Function(self.V).vector()
 
-        A_solver_custom = CustomPETScLUSolver(method=self.method)
+        A_solver_custom = PETScLUSolver(method=self.method)
         A_solver_custom.set_operator(A)
         A_solver_custom.solve(u_custom, b)
 
@@ -44,9 +63,59 @@ def testSolver(self):
 
         diff.axpy(1.0, u)
         diff.axpy(-1.0, u_custom)
-        print("Error in computed solution against fenics solver: %.3e" %(diff.norm('l2')))
+        print("solve with %s: Error in computed solution against fenics solver: %.3e" %(backend, diff.norm('l2')))
+        self.assertTrue(diff.norm('l2') < self.tol)
+
+    def compareSolveTranspose(self, backend="dolfin"):
+        """
+        Compares the solution of using custom solver against \
+            Standard solver for solving the transposed system
+        """
+        if backend == "dolfin":
+            A = dl.Matrix()
+            b = dl.Vector()
+        elif backend == "petsc":
+            A = dl.PETScMatrix()
+            b = dl.PETScVector()
+        else:
+            raise ValueError("Not a valid backend")
+
+        make_linear_system(self.V, A, b)
+        At = dl.as_backend_type(hp.Transpose(A))
+        u_fun = dl.Function(self.V)
+        u = u_fun.vector()
+        u_custom = u.copy()
+
+        diff = dl.Function(self.V).vector()
+
+        # Use the solveTranspose with custom solver 
+        A_solver_custom = PETScLUSolver(method=self.method)
+        A_solver_custom.set_operator(A)
+        A_solver_custom.solve_transpose(u_custom, b)
+
+        # a standard transpose solver 
+        At_solver = dl.PETScLUSolver(At)
+        At_solver.solve(u,b)
+
+        diff.axpy(1.0, u)
+        diff.axpy(-1.0, u_custom)
+        print("solve_tranpose with %s: Error in computed solution against fenics solver: %.3e" %(backend, diff.norm('l2')))
         self.assertTrue(diff.norm('l2') < self.tol)
 
+    def testSolveDolfin(self):
+        self.compareSolve(backend="dolfin")
+
+    def testSolveTransposeDolfin(self):
+        self.compareSolveTranspose(backend="dolfin")
+
+    def testSolvePETSc(self):
+        self.compareSolve(backend="petsc")
+
+    def testSolveTransposePETSc(self):
+        self.compareSolveTranspose(backend="petsc")
 
-if __name__ == "__main__":
+
+if __name__ == "__main__": 
     unittest.main()
+
+