Jobh/petsc solver

block/algebraic/petsc/__init__.py

@@ -18,3 +18,4 @@ def __call__(self):
 
 from .precond import *
 from .matrix import *
+from .solver import *

block/algebraic/petsc/precond.py

@@ -11,61 +11,72 @@
 from time import time
 from dolfin import info
 
-
-class precond(block_base):
-    def __init__(self, A, prectype, pdes, nullspace, prefix, options, V=None, create_cb=None, defaults={}):
-        Ad = mat(A)
-
-        if nullspace:
-            from block.block_util import isscalar
-            ns = PETSc.NullSpace()
-            if isscalar(nullspace):
-                ns.create(constant=True)
-            else:
-                ns.create(constant=False, vectors=[v.down_cast().vec() for v in nullspace])
-            try:
-                Ad.setNearNullSpace(ns)
-            except:
-                info('failed to set near null space (not supported in petsc4py version)')
-
+class petsc_base(block_base):
+    @classmethod
+    def _merge_options(cls, prefix, options, defaults):
         # The PETSc options database is global, so we create a unique prefix
         # that precisely matches the supplied options. The user can override
         # this prefix if necessary.
         if prefix is None:
             # Use class name instead of hash of defaults (we assume defaults
             # are constant for a given class)
-            prefix = self.__class__.__name__
+            prefix = cls.__name__
             if options:
                 sha = hashlib.sha256(str(tuple(sorted(options.items()))).encode())
                 prefix += sha.hexdigest()[:8]
             prefix += ':'
 
-        self.optsDB = PETSc.Options(prefix)
+        optsDB = PETSc.Options(prefix)
         params = defaults.copy()
         if options:
             params.update(options)
         for key, val in params.items():
             # Allow overriding defaults by setting key='del'. Command-line
             # options (already in optsDB) are never overwritten.
-            if val!='del' and not self.optsDB.hasName(key):
-                self.optsDB.setValue(key, val)
+            if val!='del' and not optsDB.hasName(key):
+                optsDB.setValue(key, val)
+        return prefix, optsDB
+
+    def down_cast(self):
+        return self.petsc_op
+
+class precond(petsc_base):
+    def __init__(self, A, prectype, pdes, nullspace, prefix, options, V=None, defaults={}):
+        Ad = mat(A)
+
+        if nullspace:
+            from block.block_util import isscalar
+            ns = PETSc.NullSpace()
+            if isscalar(nullspace):
+                ns.create(constant=True)
+            else:
+                ns.create(constant=False, vectors=[vec(v) for v in nullspace])
+            try:
+                Ad.setNearNullSpace(ns)
+            except:
+                info('failed to set near null space (not supported in petsc4py version)')
+
+        prefix, self.optsDB = self._merge_options(prefix=prefix, options=options, defaults=defaults)
 
         if prectype == PETSc.PC.Type.HYPRE and nullspace:
             if not all(self.optsDB.hasName(n) for n in ['pc_hypre_boomeramg_nodal_coarsen', 'pc_hypre_boomeramg_vec_interp_variant']):
                 print('Near null space is ignored by hypre UNLESS you also supply pc_hypre_boomeramg_nodal_coarsen and pc_hypre_boomeramg_vec_interp_variant')
 
         self.A = A
-        self.petsc_prec = PETSc.PC().create(V.mesh().mpi_comm() if V else None)
-        self.petsc_prec.setOptionsPrefix(prefix)
-        try:
-            self.petsc_prec.setType(prectype)
-        except PETSc.Error as e:
-            if e.ierr == 86:
-                raise ValueError(f'Unknown PETSc type "{prectype}"')
-            else:
-                raise
-        self.petsc_prec.setOperators(Ad)
-        self.petsc_prec.setFromOptions()
+        self.petsc_op = PETSc.PC().create(V.mesh().mpi_comm() if V else None)
+        self.petsc_op.setOptionsPrefix(prefix)
+
+        if prectype:
+            try:
+                self.petsc_op.setType(prectype)
+            except PETSc.Error as e:
+                if e.ierr == 86:
+                    raise ValueError(f'Unknown PETSc type "{prectype}"')
+                else:
+                    raise
+
+            self.petsc_op.setOperators(Ad)
+            self.petsc_op.setFromOptions()
 
         self.setup_time = -1.
 
@@ -74,37 +85,33 @@ def matvec(self, b):
         if self.setup_time < 0:
             T = time()
             # Create preconditioner based on the options database
-            self.petsc_prec.setUp()
+            self.petsc_op.setUp()
 
             setup_time = time()-T
-            comm = self.petsc_prec.getComm().tompi4py()
+            comm = self.petsc_op.getComm().tompi4py()
 
             setup_time = comm.allreduce(setup_time)/comm.size
             info('constructed %s preconditioner in %.2f s'%(self.__class__.__name__, setup_time))
 
             self.setup_time = setup_time
 
-        from dolfin import GenericVector
-        if not isinstance(b, GenericVector):
+        from dolfin import GenericVector, PETScVector
+        if not isinstance(b, (PETScVector, GenericVector)):
             return NotImplemented
         x = self.A.create_vec(dim=1)
         if len(x) != len(b):
             raise RuntimeError(
                 'incompatible dimensions for PETSc matvec, %d != %d'%(len(x),len(b)))
 
-        self.petsc_prec.apply(b.down_cast().vec(), x.down_cast().vec())
+        self.petsc_op.apply(vec(b), vec(x))
         return x
 
-    def down_cast(self):
-        return self.petsc_prec
-
     def __str__(self):
         return '<%s prec of %s>'%(self.__class__.__name__, str(self.A))
 
     def create_vec(self, dim=1):
         return self.A.create_vec(dim)
 
-
 class ML(precond):
     def __init__(self, A, parameters=None, pdes=1, nullspace=None, prefix=None):
         super().__init__(A, PETSc.PC.Type.ML,
@@ -128,6 +135,10 @@ def __init__(self, A, parameters=None, pdes=1, nullspace=None, prefix=None):
                          pdes=pdes, nullspace=nullspace, options=parameters, prefix=prefix)
 
 class Cholesky(precond):
+    def __init__(self, A, parameters=None, prefix=None):
+        super().__init__(A, PETSc.PC.Type.CHOLESKY, 1, None, options=parameters, prefix=prefix)
+
+class LU(precond):
     def __init__(self, A, parameters=None, prefix=None):
         super().__init__(A, PETSc.PC.Type.CHOLESKY,
                          pdes=1, nullspace=None, options=parameters, prefix=prefix)
@@ -139,12 +150,21 @@ def __init__(self, A, parameters=None, prefix=None, defaults={}):
                          defaults=defaults)
 
 
+class UMFPACK_LU(LU):
+    def __init__(self, A, parameters=None, prefix=None):
+        super().__init__(A,
+                         parameters=parameters, prefix=prefix,
+                         defaults={
+                             'pc_factor_mat_solver_type': 'umfpack',
+                         })
+
+
 class MUMPS_LU(LU):
     def __init__(self, A, parameters=None, prefix=None):
         super().__init__(A,
                          parameters=parameters, prefix=prefix,
                          defaults={
-                             'pc_factor_mat_solver_package': 'mumps',
+                             'pc_factor_mat_solver_type': 'mumps',
                          })
 
 
@@ -153,9 +173,9 @@ def __init__(self, A, parameters=None, prefix=None):
         super().__init__(A,
                          parameters=parameters, prefix=prefix,
                          defaults={
-                             'pc_factor_mat_solver_package': 'superlu',
+                             'pc_factor_mat_solver_type': 'superlu',
                          })
-        self.petsc_prec.setFactorPivot(1E-16)
+        self.petsc_op.setFactorPivot(1E-16)
 
 class AMG(precond):
     """
@@ -298,17 +318,19 @@ def __init__(self, A, V, ams_zero_beta_poisson=False, prefix=None):
                          defaults={
                              'pc_hypre_type': 'ams',
                          })
-        self.petsc_prec.setHYPREDiscreteGradient(mat(G))
+        self.petsc_op.setHYPREDiscreteGradient(mat(G))
 
         # Constants
         constants = [vec(df.interpolate(df.Constant(c), V).vector())
                      for c in np.eye(mesh.geometry().dim())]
-        self.petsc_prec.setHYPRESetEdgeConstantVectors(*constants)
+        self.petsc_op.setHYPRESetEdgeConstantVectors(*constants)
 
         # NOTE: signal that we don't have lower order term
-        ams_zero_beta_poisson and self.petsc_prec.setHYPRESetBetaPoissonMatrix(None)
+        ams_zero_beta_poisson and self.petsc_op.setHYPRESetBetaPoissonMatrix(None)
 
+        self.setup_time = -1
 
+
 class HypreADS(precond):
     '''
     AMG auxiliary space preconditioner for H(div) problems in 3d. The bilinear 
@@ -373,20 +395,21 @@ def __init__(self, A, V, prefix=None):
         assert V.mesh().geometry().dim() == 3
         assert V.ufl_element().family() == 'Raviart-Thomas'
         assert V.ufl_element().degree() == 1
-
         supports_mpi(False, 'ADS curl matrix needs to be fixed for parallel', mat(A).comm.size) # FIXME
 
         mesh = V.mesh()
         C = HypreADS.discrete_curl(mesh)
         G = HypreADS.discrete_gradient(mesh)
         xyz = HypreADS.coordinates(mesh)
 
-        super().__init__(A, PETSc.PC.Type.HYPRE, V=V,
-                         pdes=None, nullspace=None, options=None, prefix=prefix,
-                         defaults={
-                             'pc_hypre_type': 'ads',
-                         })
+        super().__init__(A, '', V=V,
+                         pdes=None, nullspace=None, options=None, prefix=prefix)
+        self.petsc_op.setType('hypre')        
+        self.petsc_op.setHYPREType('ads')
+        self.petsc_op.setOperators(mat(A))
         # Attach auxiliary operators
-        self.petsc_prec.setHYPREDiscreteGradient(G)
-        self.petsc_prec.setHYPREDiscreteCurl(G)
-        self.petsc_prec.setCoordinates(xyz)
+        self.petsc_op.setHYPREDiscreteGradient(G)
+        self.petsc_op.setHYPREDiscreteCurl(G)
+        self.petsc_op.setCoordinates(xyz)
+        self.petsc_op.setUp()
+