diff --git a/soupy/solver/PETScLUSolver.py b/soupy/solver/PETScLUSolver.py new file mode 100644 index 0000000..809633a --- /dev/null +++ b/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("") diff --git a/soupy/solver/__init__.py b/soupy/solver/__init__.py index 9fb5610..7e2687b 100644 --- a/soupy/solver/__init__.py +++ b/soupy/solver/__init__.py @@ -15,4 +15,4 @@ from .nonlinearVariationalSolver import NonlinearVariationalSolver -from .customPETScLUSolver import CustomPETScLUSolver +from .PETScLUSolver import PETScLUSolver diff --git a/soupy/solver/customPETScLUSolver.py b/soupy/solver/customPETScLUSolver.py deleted file mode 100644 index 7ecdbf7..0000000 --- a/soupy/solver/customPETScLUSolver.py +++ /dev/null @@ -1,73 +0,0 @@ -import dolfin as dl -from petsc4py import PETSc -from mpi4py import MPI - - -class CustomPETScLUSolver: - def __init__(self, mpi_comm=MPI.COMM_WORLD, method="mumps", reuse=False): - self.mpi_comm = mpi_comm - self.ksp = PETSc.KSP().create(comm=self.mpi_comm) - self.ksp.setType('preonly') - self.pc = self.ksp.getPC() - self.pc.setType('lu') - opts = PETSc.Options() - opts['pc_factor_mat_solver_type'] = method - self.pc.setFromOptions() - if reuse: - self.pc.setReusePreconditioner(True) - self.verbose = False - - def set_operator(self, A): - A_mat = dl.as_backend_type(A).mat() - self.ksp.setOperators(A_mat, A_mat) - - - def setReuseFactorization(self, flag): - self.pc.setReusePreconditioner(flag) - - - def solve(self, x, b): - if self.verbose: - print("Calling custom solve") - b_vec = dl.as_backend_type(b).vec() - x_petsc = dl.as_backend_type(x) - x_vec = x_petsc.vec() - self.ksp.solve(b_vec, x_vec) - x.set_local(x_petsc.get_local()) - x.apply("") - -if __name__ == "__main__": - import matplotlib.pyplot as plt - mesh = dl.UnitSquareMesh(20,20) - V = dl.FunctionSpace(mesh, "CG", 2) - u_trial = dl.TrialFunction(V) - u_test = dl.TestFunction(V) - u_fun = dl.Function(V) - u = u_fun.vector() - - 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") - - A = dl.PETScMatrix() - b = dl.PETScVector() - dl.assemble_system(a_form, b_form, bcs=bcs, A_tensor=A, b_tensor=b) - - A_solver = dl.PETScLUSolver(A) - A_solver.solve(u,b) - - plt.figure() - p = dl.plot(u_fun) - plt.colorbar(p) - plt.title("Fenics built in") - - A_solver_custom = CustomPETScLuSolver() - A_solver_custom.set_operator(A) - A_solver_custom.solve(u, b) - - plt.figure() - p = dl.plot(u_fun) - plt.colorbar(p) - plt.title("Custom") - plt.show() diff --git a/soupy/test/test_customPETScLUSolver.py b/soupy/test/test_customPETScLUSolver.py index 4de3651..6bae2f9 100644 --- a/soupy/test/test_customPETScLUSolver.py +++ b/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() + +