diff --git a/block/algebraic/petsc/__init__.py b/block/algebraic/petsc/__init__.py index 98cbfce..3419ebb 100644 --- a/block/algebraic/petsc/__init__.py +++ b/block/algebraic/petsc/__init__.py @@ -18,3 +18,4 @@ def __call__(self): from .precond import * from .matrix import * +from .solver import * diff --git a/block/algebraic/petsc/precond.py b/block/algebraic/petsc/precond.py index 873cbc7..ae1687c 100644 --- a/block/algebraic/petsc/precond.py +++ b/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,7 +395,6 @@ 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() @@ -381,12 +402,14 @@ def __init__(self, A, V, prefix=None): 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() + diff --git a/block/algebraic/petsc/solver.py b/block/algebraic/petsc/solver.py new file mode 100644 index 0000000..079e11a --- /dev/null +++ b/block/algebraic/petsc/solver.py @@ -0,0 +1,195 @@ +from block import block_mat, block_vec, supports_mpi +from block.block_base import block_base, block_container +from block.object_pool import vec_pool +from dolfin import GenericVector, Matrix, PETScVector, PETScMatrix +import numpy as np +from petsc4py import PETSc +from .precond import petsc_base, mat as single_mat, vec as single_vec + +# VecNest is a perfect fit for block_vec, but I can't get it to work (gives +# error code 86), and I don't know how to get proper error messages or traces +# from PETSc. When set to False, it works but +# - is very inefficient (lots of copying of vectors) +# - doesn't work in parallel +USE_VECNEST=False + +def mat(A, _sizes=None): + if isinstance(A, block_mat): + # We need sizes upfront in case there are matrix-free blocks, creating + # a vector to find the size is wasteful but easy + row_sz = [(v.local_size(), v.size()) for v in A.create_vec(dim=1)] + col_sz = [(v.local_size(), v.size()) for v in A.create_vec(dim=0)] + + m = PETSc.Mat().createNest([[mat(m, (row_sz[r], col_sz[c])) for (c,m) in enumerate(row)] + for (r,row) in enumerate(A.blocks)]) + if USE_VECNEST: + m.setVecType(PETSc.Vec.Type.NEST) + + return m + + elif isinstance(A, (PETScMatrix, Matrix)): + return single_mat(A) + + elif isinstance(A, block_base) or np.isscalar(A): + if not _sizes: + raise ValueError('Cannot create unsized PETSc matrix') + + if A == 0: + return None + + Ad = PETSc.Mat().createPython(_sizes) + Ad.setPythonContext(petsc_py_wrapper(A)) + Ad.setUp() + return Ad + else: + try: + row_vec, col_vec = A.create_vec(dim=1), A.create_vec(dim=0) + _sizes = (row_vec.local_size(), col_vec.local_size()) + + Ad = PETSc.Mat().createPython(_sizes) + Ad.setPythonContext(petsc_py_wrapper(A)) + Ad.setUp() + + return Ad + except: + raise TypeError(str(type(A))) + +def vec(v): + if isinstance(v, block_vec): + if USE_VECNEST: + return PETSc.Vec().createNest([single_vec(x) for x in v]) + else: + vecs = [vec(x) for x in v] + sizes = [x.getSizes() for x in vecs] + vs = PETSc.Vec().create(comm=vecs[0].comm) + vs.setType(PETSc.Vec.Type.STANDARD) + vs.setSizes(sum(s[1] for s in sizes)) + arr = vs.getArray() + i0 = 0 + for vv in vecs: + arr2 = vv.getArray() + arr[i0:i0+len(arr2)] = arr2 + i0 += len(arr2) + return vs + elif isinstance(v, (PETScVector, GenericVector)): + return single_vec(v) + else: + raise TypeError(str(type(v))) + +def Vec(v, creator): + if isinstance(v, PETSc.Vec): + if USE_VECNEST: + if v.type == PETSc.Vec.Type.NEST: + return block_vec([Vec(x) for x in v]) + else: + return PETScVector(v) + else: + try: + ret = creator.create_vec(dim=1) + except: + pass + + if isinstance(ret, block_vec): + arr = v.getArray() + i0 = 0 + for vv in ret: + # vv.set_local(arr[i0:i0+len(vv)]) + vv *= 0. + vv.axpy(1., PETScVector(PETSc.Vec().createWithArray(arr[i0:i0+len(vv)]))) + i0 += len(vv) + return ret + else: + return PETScVector(v) + else: + raise TypeError(str(type(v))) + +class petsc_py_wrapper: + # Python context for cbc.block actions (PC.Type.PYTHON, Mat.Type.PYTHON) + def __init__(self, A): + self.A = A + + def mult(self, mat, x, y): + new_y = vec(self.A * Vec(x, self.A)) + # y.aypx(0, new_y) + y *= 0 + y += new_y + + def apply(self, pc, x, y): + return self.mult(None, x, y) + + def setUp(self, pc): + pass + +class petsc_solver(petsc_base): + def __init__(self, A, precond, V, ksp_type, prefix, options, defaults={}, nullspace=None): + supports_mpi(False, 'PETSc solver interface currently does not work in parallel') + + self.A = A + self.Ad = mat(A) + + prefix, self.optsDB = self._merge_options(prefix=prefix, options=options, defaults=defaults) + + self.petsc_op = PETSc.KSP().create(V.mesh().mpi_comm() if V else None) + self.petsc_op.setOptionsPrefix(prefix) + if ksp_type is not None: + self.petsc_op.setType(ksp_type) + + self.petsc_op.setConvergenceHistory() # Record residuals + self.residuals = [] + + self.petsc_op.setComputeEigenvalues(True) + + if nullspace is not None: + ns = PETSc.NullSpace().create(vectors=[vec(fi) for fi in nullspace]) + self.Ad.setNearNullSpace(ns) + self.petsc_op.setOperators(self.Ad) + + self.petsc_op.setFromOptions() + + if precond is not None: + pc = self.petsc_op.getPC() + pc.setType(PETSc.PC.Type.PYTHON) + pc.setPythonContext(petsc_py_wrapper(precond)) + + def __str__(self): + return '<%s ksp of %s>'%(self.__class__.__name__, str(self.A)) + + def matvec(self, b): + self.petsc_op.setUp() + x = self.Ad.createVecLeft() + self.petsc_op.solve(vec(b), x) + + self.residuals = self.petsc_op.getConvergenceHistory() + + return Vec(x, self.A) + + @vec_pool + def create_vec(self, dim=1): + from dolfin import PETScVector + if dim == 0: + m = self.Ad.createVecRight() + elif dim == 1: + m = self.Ad.createVecLeft() + else: + raise ValueError('dim must be 0 or 1') + return PETScVector(m) + + @property + def iterations(self): + return len(self.residuals)-1 + + def eigenvalue_estimates(self): + if self.petsc_op.getComputeEigenvalues(): + return self.petsc_op.computeEigenvalues() + + if self.petsc_op.getComputeSingularValues(): + return self.petsc_op.computeExtremeSingularValues() + + return np.array([]) + +class KSP(petsc_solver): + def __init__(self, A, precond=None, prefix=None, nullspace=None, **parameters): + super().__init__(A, precond=precond, V=None, ksp_type=None, prefix=prefix, options=parameters, + nullspace=nullspace, + defaults={ + }) diff --git a/block/block_mat.py b/block/block_mat.py index 1c55d1d..8dc6213 100644 --- a/block/block_mat.py +++ b/block/block_mat.py @@ -2,6 +2,8 @@ from builtins import range from .block_base import block_container from .block_vec import block_vec +import numpy as np + class block_mat(block_container): """Block of matrices or other operators. Empty blocks doesn't need to be set @@ -18,7 +20,7 @@ def __init__(self, m, n=None, blocks=0): block_container.__init__(self, (m,n), blocks) def matvec(self, x): - from dolfin import GenericVector, Matrix, PETScMatrix + from dolfin import GenericVector, Matrix, PETScMatrix, PETScVector from .block_util import isscalar import numpy m,n = self.blocks.shape @@ -45,7 +47,16 @@ def matvec(self, x): # Do the block multiply if isinstance(self[i,j], PETScMatrix): z = self[i,j].create_vec(dim=0) - self[i,j].mult(x[j], z) + + if isinstance(x[j], (GenericVector, PETScVector)): + self[i,j].mult(x[j], z) + else: + xx = self.create_vec() + offsets = np.cumsum([0] + [xi.size() for xi in xx]) + xj = xx[j] + + xj.set_local(x.get_local()[offsets[j]:offsets[j+1]]) + self[i, j].mult(xj, z) else: z = self[i,j] * x[j] if isinstance(z, type(NotImplemented)): return NotImplemented diff --git a/block/block_vec.py b/block/block_vec.py index 8b2df96..9b15148 100644 --- a/block/block_vec.py +++ b/block/block_vec.py @@ -43,8 +43,12 @@ def allocate(self, template, dim=None, alternative_templates=[]): except Exception: pass if not isinstance(self[i], GenericVector): + from IPython import embed + embed() + raise ValueError( f"Can't allocate vector - no usable template for block {i}.\n" + f'{self[i]}' "Consider calling something like bb.allocate([V, Q]) to initialise the block_vec." ) self[i].zero() diff --git a/data/regression/petsc.Sigma.pickle b/data/regression/petsc.Sigma.pickle new file mode 100644 index 0000000..9e62a78 Binary files /dev/null and b/data/regression/petsc.Sigma.pickle differ diff --git a/data/regression/petsc.U.pickle b/data/regression/petsc.U.pickle new file mode 100644 index 0000000..994128d Binary files /dev/null and b/data/regression/petsc.U.pickle differ diff --git a/demo/petsc.py b/demo/petsc.py new file mode 100644 index 0000000..f3ab95c --- /dev/null +++ b/demo/petsc.py @@ -0,0 +1,109 @@ +from block import * +from block.algebraic.hazmath import AMG as hazAMG +from block.algebraic.petsc import KSP, AMG +from block.dolfin_util import BoxBoundary +from block.iterative import Richardson +from dolfin import * + +######################### +# Basic Poisson (poisson.py) +######################### + +mesh = UnitSquareMesh(32,32) +V = FunctionSpace(mesh, "CG", 1) + +f = Expression("sin(pi*x[0])", degree=2) +u, v = TrialFunction(V), TestFunction(V) + +a = u*v*dx + dot(grad(u), grad(v))*dx +L = f*v*dx + +A = assemble(a) +b = assemble(L) + +u = Function(V).vector() +solve(A, u, b) + +######################### +# Solve as single-matrix +######################### + +Apre = AMG(A) +Ainv = KSP(A, ksp_type='cg', pc_type='ilu') +# Alternatively: +#Ainv = KSP(A, ksp_type='cg', pc_type='hypre', pc_hypre_type='boomeramg') +x = Ainv*b + +check_expected('x', x, expected=u, show=True) + +######################### +# Solve a (trivial) block-matrix, using KSP Ainv as preconditioner +######################### + +AA = block_mat([[A]]) +bb = block_vec([b]) +AApre = block_mat([[Ainv]]) +AAinv = KSP(AA, precond=AApre, ksp_type='cg') + +xx = AAinv*bb + +check_expected('x', xx[0], expected=u, show=True) + +######################### +# Solve a nontrivial block-matrix (mixedpoisson.py) +######################### + +BDM = FunctionSpace(mesh, "BDM", 1) +DG = FunctionSpace(mesh, "DG", 0) + +tau, sigma = TestFunction(BDM), TrialFunction(BDM) +v, u = TestFunction(DG), TrialFunction(DG) + +class BoundarySource(UserExpression): + def __init__(self, mesh): + super().__init__(self) + self.mesh = mesh + def eval_cell(self, values, x, ufc_cell): + cell = Cell(self.mesh, ufc_cell.index) + n = cell.normal(ufc_cell.local_facet) + g = sin(5*x[0]) + values[0] = g*n[0] + values[1] = g*n[1] + def value_shape(self): + return (2,) + +G = BoundarySource(mesh) +boundary = BoxBoundary(mesh).ns +bcs_BDM = [DirichletBC(BDM, G, boundary)] +bcs = block_bc([bcs_BDM, None], symmetric=True) + +f = Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=4) + +a11 = dot(sigma, tau) * dx +a12 = div(tau) * u * dx +a21 = div(sigma) * v *dx +L2 = - f * v * dx + +AA = block_assemble([[a11, a12], + [a21, 0 ]]) + +bb = block_assemble([0, L2]) + +bcs.apply(AA).apply(bb) + +# Extract the individual submatrices +[[A, B], + [C, _]] = AA + +Ap = hazAMG(A) +Lp = Richardson(C*B, precond=0.5, iter=40, name='L^', show=0) + +AAp = block_mat([[Ap, 0], + [0, Lp]]) + +AAinv = KSP(AA, precond=AAp, ksp_type='minres') +Sigma, U = AAinv * bb + +print('expected (from mixedpoisson.py): Sigma=1.2138, U=6.7166') +check_expected('Sigma', Sigma, show=True) +check_expected('U', U, show=True)