From 5d13f9e5f5a302b07c1e62f1372985d1eadf9523 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 10 Jan 2023 13:22:21 +0000 Subject: [PATCH 01/14] Add FacetSplitPC --- firedrake/preconditioners/facet_split.py | 252 +++++++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 firedrake/preconditioners/facet_split.py diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py new file mode 100644 index 0000000000..a7d10303c1 --- /dev/null +++ b/firedrake/preconditioners/facet_split.py @@ -0,0 +1,252 @@ +from firedrake.petsc import PETSc +from firedrake.preconditioners.base import PCBase +import firedrake.dmhooks as dmhooks +import numpy + + +__all__ = ['FacetSplitPC'] + + +def get_permutation_map(V, W): + from firedrake.preconditioners.fdm import glonum_fun, restricted_dofs + bsize = V.value_size + V_indices = None + V_map, nel = glonum_fun(V.cell_node_map(), bsize=bsize) + perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) + perm.fill(-1) + + offset = 0 + for Wsub in W: + W_indices = None + W_map, nel = glonum_fun(Wsub.cell_node_map(), bsize=bsize) + rdofs = restricted_dofs(Wsub.finat_element, V.finat_element) + + for e in range(nel): + V_indices = V_map(e, result=V_indices) + W_indices = W_map(e, result=W_indices) + perm[V_indices[rdofs]] = W_indices + offset + + offset += Wsub.dof_dset.set.size * bsize + del rdofs + del W_indices + del V_indices + perm = V.dof_dset.lgmap.apply(perm, result=perm) + own = V.dof_dset.set.size * bsize + return perm[:own] + + +def get_permutation_project(V, W): + from firedrake import Function, split + ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() + start, end = ownership_ranges[V.comm.rank:V.comm.rank+2] + v = Function(V) + w = Function(W) + with w.dat.vec_wo as wvec: + wvec.setArray(numpy.linspace(0.0E0, 1.0E0, end-start, dtype=PETSc.RealType)) + w_expr = sum(split(w)) + try: + v.interpolate(w_expr) + except NotImplementedError: + rtol = 1.0 / max(numpy.diff(ownership_ranges))**2 + v.project(w_expr, solver_parameters={ + "mat_type": "matfree", + "ksp_type": "cg", + "ksp_atol": 0, + "ksp_rtol": rtol, + "pc_type": "jacobi", }) + return numpy.rint((end-start-1)*v.dat.data_ro.reshape((-1,))+start).astype(PETSc.IntType) + + +class FacetSplitPC(PCBase): + + needs_python_pmat = False + _prefix = "facet_" + + _permutation_cache = {} + + def get_permutation(self, V, W): + from mpi4py import MPI + key = (V, W) + if key not in self._permutation_cache: + indices = get_permutation_map(V, W) + if V.comm.allreduce(numpy.all(indices[:-1] <= indices[1:]), MPI.PROD): + self._permutation_cache[key] = None + else: + self._permutation_cache[key] = indices + return self._permutation_cache[key] + + def initialize(self, pc): + + from ufl import (InteriorElement, FacetElement, MixedElement, TensorElement, VectorElement, + FiniteElement, EnrichedElement, TensorProductElement, TensorProductCell, HCurl) + from firedrake import FunctionSpace, TestFunctions, TrialFunctions + from firedrake.assemble import allocate_matrix, TwoFormAssembler + from firedrake.solving_utils import _SNESContext + from functools import partial + + _, P = pc.getOperators() + appctx = self.get_appctx(pc) + fcp = appctx.get("form_compiler_parameters") + + if pc.getType() != "python": + raise ValueError("Expecting PC type python") + + ctx = dmhooks.get_appctx(pc.getDM()) + if ctx is None: + raise ValueError("No context found.") + if not isinstance(ctx, _SNESContext): + raise ValueError("Don't know how to get form from %r", ctx) + + prefix = pc.getOptionsPrefix() + options_prefix = prefix + self._prefix + options = PETSc.Options(options_prefix) + mat_type = options.getString("mat_type", "submatrix") + + problem = ctx._problem + a = problem.Jp or problem.J + V = a.arguments()[-1].function_space() + assert len(V) == 1 + + # W = V_interior * V_facet + scalar_element = V.ufl_element() + if isinstance(scalar_element, (TensorElement, VectorElement)): + scalar_element = scalar_element._sub_element + tensorize = lambda e, shape: TensorElement(e, shape=shape) if shape else e + + def get_facet_element(e): + cell = e.cell() + if e.sobolev_space() == HCurl and isinstance(cell, TensorProductCell): + sub_cells = cell.sub_cells() + degree = max(e.degree()) + variant = e.variant() + Qc_elt = FiniteElement("Q", sub_cells[0], degree, variant=variant) + Qd_elt = FiniteElement("RTCE", sub_cells[0], degree, variant=variant) + Id_elt = FiniteElement("DG", sub_cells[1], degree - 1, variant=variant) + Ic_elt = FiniteElement("CG", sub_cells[1], degree, variant=variant) + return EnrichedElement(HCurl(TensorProductElement(FacetElement(Qc_elt), Id_elt, cell=cell)), + HCurl(TensorProductElement(Qd_elt, FacetElement(Ic_elt), cell=cell)), + HCurl(TensorProductElement(FacetElement(Qd_elt), InteriorElement(Ic_elt), cell=cell))) + return FacetElement(e) + + elements = [tensorize(restriction(scalar_element), V.shape) for restriction in (InteriorElement, get_facet_element)] + + W = FunctionSpace(V.mesh(), MixedElement(elements)) + if W.dim() != V.dim(): + raise ValueError("Dimensions of the original and decomposed spaces do not match") + + self.perm = None + self.iperm = None + indices = self.get_permutation(V, W) + if indices is not None: + self.perm = PETSc.IS().createGeneral(indices, comm=V.comm) + self.iperm = self.perm.invertPermutation() + + mixed_operator = a(sum(TestFunctions(W)), sum(TrialFunctions(W)), coefficients={}) + mixed_bcs = tuple(bc.reconstruct(V=W[-1], g=0) for bc in problem.bcs) + + def _permute_nullspace(nsp): + if nsp is None or self.iperm is None: + return nsp + vecs = [vec.duplicate() for vec in nsp.getVecs()] + for vec in vecs: + vec.permute(self.iperm) + return PETSc.NullSpace().create(constant=nsp.constant, vectors=vecs, comm=nsp.comm) + + if mat_type != "submatrix": + self.mixed_op = allocate_matrix(mixed_operator, + bcs=mixed_bcs, + form_compiler_parameters=fcp, + mat_type=mat_type, + options_prefix=options_prefix) + self._assemble_mixed_op = TwoFormAssembler(mixed_operator, tensor=self.mixed_op, + form_compiler_parameters=fcp, + bcs=mixed_bcs).assemble + self._assemble_mixed_op() + mixed_opmat = self.mixed_op.petscmat + + if False: + # FIXME + mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace())) + mixed_opmat.setNearNullSpace(_permute_nullspace(P.getNearNullSpace())) + mixed_opmat.setTransposeNullSpace(_permute_nullspace(P.getTransposeNullSpace())) + + elif self.perm: + self._permute_op = partial(PETSc.Mat().createSubMatrixVirtual, P, self.iperm, self.iperm) + mixed_opmat = self._permute_op() + else: + mixed_opmat = P + + # Internally, we just set up a PC object that the user can configure + # however from the PETSc command line. Since PC allows the user to specify + # a KSP, we can do iterative by -facet_pc_type ksp. + scpc = PETSc.PC().create(comm=pc.comm) + scpc.incrementTabLevel(1, parent=pc) + + # We set a DM and an appropriate SNESContext on the constructed PC so one + # can do e.g. fieldsplit. + mixed_dm = W.dm + self._dm = mixed_dm + + # Create new appctx + self._ctx_ref = self.new_snes_ctx(pc, + mixed_operator, + mixed_bcs, + mat_type, + fcp, + options_prefix=options_prefix) + + scpc.setDM(mixed_dm) + scpc.setOptionsPrefix(options_prefix) + scpc.setOperators(A=mixed_opmat, P=mixed_opmat) + with dmhooks.add_hooks(mixed_dm, self, appctx=self._ctx_ref, save=False): + scpc.setFromOptions() + self.pc = scpc + + def update(self, pc): + if hasattr(self, "mixed_op"): + self._assemble_mixed_op() + elif hasattr(self, "_permute_op"): + for mat in self.pc.getOperators(): + mat.destroy() + P = self._permute_op() + self.pc.setOperators(A=P, P=P) + self.pc.setUp() + + def apply(self, pc, x, y): + if self.perm: + x.permute(self.iperm) + dm = self._dm + with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): + self.pc.apply(x, y) + if self.perm: + x.permute(self.perm) + y.permute(self.perm) + + def applyTranspose(self, pc, x, y): + if self.perm: + x.permute(self.iperm) + dm = self._dm + with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): + self.pc.applyTranspose(x, y) + if self.perm: + x.permute(self.perm) + y.permute(self.perm) + + def view(self, pc, viewer=None): + super(FacetSplitPC, self).view(pc, viewer) + if hasattr(self, "pc"): + viewer.printfASCII("PC using interior-facet decomposition\n") + self.pc.view(viewer) + + def destroy(self, pc): + if hasattr(self, "pc"): + if hasattr(self, "_permute_op"): + for mat in self.pc.getOperators(): + mat.destroy() + self.pc.destroy() + if hasattr(self, "iperm"): + if self.iperm: + self.iperm.destroy() + if hasattr(self, "perm"): + if self.perm: + self.perm.destroy() From 46e75de2faf87c8c404bd22468d244a54fd0b2ff Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 11 Jan 2023 16:23:15 +0000 Subject: [PATCH 02/14] get rid of python loop over cells and use general restriction onto facets --- firedrake/preconditioners/__init__.py | 1 + firedrake/preconditioners/facet_split.py | 120 +++++++++++++++-------- 2 files changed, 82 insertions(+), 39 deletions(-) diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py index 3f67ed1718..e63a3c38f1 100644 --- a/firedrake/preconditioners/__init__.py +++ b/firedrake/preconditioners/__init__.py @@ -10,3 +10,4 @@ from firedrake.preconditioners.hypre_ams import * # noqa: F401 from firedrake.preconditioners.hypre_ads import * # noqa: F401 from firedrake.preconditioners.fdm import * # noqa: F401 +from firedrake.preconditioners.facet_split import * # noqa: F401 diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index a7d10303c1..8877cc7f1e 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -7,31 +7,86 @@ __all__ = ['FacetSplitPC'] +def split_dofs(elem): + entity_dofs = elem.entity_dofs() + ndim = elem.cell.get_spatial_dimension() + edofs = [[] for k in range(ndim+1)] + for key in entity_dofs: + vals = entity_dofs[key] + edim = key + try: + edim = sum(edim) + except TypeError: + pass + split = numpy.arange(0, len(vals)+1, 2**(ndim-edim)) + for r in range(len(split)-1): + v = sum([vals[k] for k in range(*split[r:r+2])], []) + edofs[edim].extend(sorted(v)) + + return numpy.array(edofs[-1], dtype=PETSc.IntType), numpy.array(sum(reversed(edofs[:-1]), []), dtype=PETSc.IntType) + + +def restricted_dofs(celem, felem): + """ + find which DOFs from felem are on celem + :arg celem: the restricted :class:`finat.FiniteElement` + :arg felem: the unrestricted :class:`finat.FiniteElement` + :returns: :class:`numpy.array` with indices of felem that correspond to celem + """ + csplit = split_dofs(celem) + fsplit = split_dofs(felem) + if len(csplit[0]) and len(csplit[1]): + csplit = [numpy.concatenate(csplit)] + fsplit = [numpy.concatenate(fsplit)] + + k = len(csplit[0]) == 0 + if len(csplit[k]) != len(fsplit[k]): + raise ValueError("Finite elements have different DOFs") + perm = numpy.empty_like(csplit[k]) + perm[csplit[k]] = numpy.arange(len(perm), dtype=perm.dtype) + return fsplit[k][perm] + + def get_permutation_map(V, W): - from firedrake.preconditioners.fdm import glonum_fun, restricted_dofs + from firedrake import Function + from pyop2 import op2, PermutedMap + bsize = V.value_size - V_indices = None - V_map, nel = glonum_fun(V.cell_node_map(), bsize=bsize) perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) perm.fill(-1) + v = Function(V, dtype=PETSc.IntType, val=perm) + w = Function(W, dtype=PETSc.IntType) offset = 0 - for Wsub in W: - W_indices = None - W_map, nel = glonum_fun(Wsub.cell_node_map(), bsize=bsize) - rdofs = restricted_dofs(Wsub.finat_element, V.finat_element) - - for e in range(nel): - V_indices = V_map(e, result=V_indices) - W_indices = W_map(e, result=W_indices) - perm[V_indices[rdofs]] = W_indices + offset - - offset += Wsub.dof_dset.set.size * bsize - del rdofs - del W_indices - del V_indices - perm = V.dof_dset.lgmap.apply(perm, result=perm) + for wdata, Wsub in zip(w.dat.data, W): + own = Wsub.dof_dset.set.size * bsize + wdata[:own] = numpy.arange(offset, offset+own, dtype=PETSc.IntType) + offset += own + + idofs = W[0].finat_element.space_dimension() * bsize + fdofs = W[1].finat_element.space_dimension() * bsize + + eperm = numpy.concatenate([restricted_dofs(Wsub.finat_element, V.finat_element) for Wsub in W]) + pmap = PermutedMap(V.cell_node_map(), eperm) + + kernel_code = f""" + void permutation(PetscInt *restrict x, + const PetscInt *restrict xi, + const PetscInt *restrict xf){{ + + for(PetscInt i=0; i<{idofs}; i++) x[i] = xi[i]; + for(PetscInt i=0; i<{fdofs}; i++) x[i+{idofs}] = xf[i]; + return; + }} + """ + kernel = op2.Kernel(kernel_code, "permutation", requires_zeroed_output_arguments=False) + op2.par_loop(kernel, v.cell_set, + v.dat(op2.WRITE, pmap), + w.dat[0](op2.READ, W[0].cell_node_map()), + w.dat[1](op2.READ, W[1].cell_node_map())) + own = V.dof_dset.set.size * bsize + perm = V.dof_dset.lgmap.apply(perm, result=perm) return perm[:own] @@ -58,6 +113,11 @@ def get_permutation_project(V, W): class FacetSplitPC(PCBase): + """ A preconditioner that splits a function into interior and facet DOFs. + + Composability with fieldsplit can be done through the PETSc solver option: + -facet_pc_type fieldsplit + """ needs_python_pmat = False _prefix = "facet_" @@ -77,8 +137,7 @@ def get_permutation(self, V, W): def initialize(self, pc): - from ufl import (InteriorElement, FacetElement, MixedElement, TensorElement, VectorElement, - FiniteElement, EnrichedElement, TensorProductElement, TensorProductCell, HCurl) + from ufl import InteriorElement, FacetElement, MixedElement, TensorElement, VectorElement from firedrake import FunctionSpace, TestFunctions, TrialFunctions from firedrake.assemble import allocate_matrix, TwoFormAssembler from firedrake.solving_utils import _SNESContext @@ -112,27 +171,10 @@ def initialize(self, pc): if isinstance(scalar_element, (TensorElement, VectorElement)): scalar_element = scalar_element._sub_element tensorize = lambda e, shape: TensorElement(e, shape=shape) if shape else e - - def get_facet_element(e): - cell = e.cell() - if e.sobolev_space() == HCurl and isinstance(cell, TensorProductCell): - sub_cells = cell.sub_cells() - degree = max(e.degree()) - variant = e.variant() - Qc_elt = FiniteElement("Q", sub_cells[0], degree, variant=variant) - Qd_elt = FiniteElement("RTCE", sub_cells[0], degree, variant=variant) - Id_elt = FiniteElement("DG", sub_cells[1], degree - 1, variant=variant) - Ic_elt = FiniteElement("CG", sub_cells[1], degree, variant=variant) - return EnrichedElement(HCurl(TensorProductElement(FacetElement(Qc_elt), Id_elt, cell=cell)), - HCurl(TensorProductElement(Qd_elt, FacetElement(Ic_elt), cell=cell)), - HCurl(TensorProductElement(FacetElement(Qd_elt), InteriorElement(Ic_elt), cell=cell))) - return FacetElement(e) - - elements = [tensorize(restriction(scalar_element), V.shape) for restriction in (InteriorElement, get_facet_element)] + elements = [tensorize(restriction(scalar_element), V.shape) for restriction in (InteriorElement, FacetElement)] W = FunctionSpace(V.mesh(), MixedElement(elements)) - if W.dim() != V.dim(): - raise ValueError("Dimensions of the original and decomposed spaces do not match") + assert W.dim() == V.dim(), "Dimensions of the original and decomposed spaces do not match" self.perm = None self.iperm = None From 0f8e199f17df23850bfbe3b2151954a75bc2cfe4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 16 Jan 2023 11:39:34 +0000 Subject: [PATCH 03/14] add a test --- firedrake/preconditioners/facet_split.py | 55 ++++++++++++++---------- tests/regression/test_facet_split.py | 38 ++++++++++++++++ 2 files changed, 70 insertions(+), 23 deletions(-) create mode 100644 tests/regression/test_facet_split.py diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 8877cc7f1e..491c0118a1 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -137,7 +137,7 @@ def get_permutation(self, V, W): def initialize(self, pc): - from ufl import InteriorElement, FacetElement, MixedElement, TensorElement, VectorElement + from ufl import RestrictedElement, MixedElement, TensorElement, VectorElement from firedrake import FunctionSpace, TestFunctions, TrialFunctions from firedrake.assemble import allocate_matrix, TwoFormAssembler from firedrake.solving_utils import _SNESContext @@ -161,21 +161,33 @@ def initialize(self, pc): options = PETSc.Options(options_prefix) mat_type = options.getString("mat_type", "submatrix") - problem = ctx._problem - a = problem.Jp or problem.J - V = a.arguments()[-1].function_space() - assert len(V) == 1 + if P.getType() == "python": + ictx = P.getPythonContext() + a = ictx.a + bcs = tuple(ictx.row_bcs) + else: + problem = ctx._problem + a = problem.Jp or problem.J + bcs = tuple(problem.bcs) - # W = V_interior * V_facet - scalar_element = V.ufl_element() - if isinstance(scalar_element, (TensorElement, VectorElement)): - scalar_element = scalar_element._sub_element - tensorize = lambda e, shape: TensorElement(e, shape=shape) if shape else e - elements = [tensorize(restriction(scalar_element), V.shape) for restriction in (InteriorElement, FacetElement)] + V = a.arguments()[-1].function_space() + assert len(V) == 1, "Interior-facet decomposition of mixed elements is not supported" + + # W = V[interior] * V[facet] + def restrict(ele, restriction_domain): + if isinstance(ele, VectorElement): + return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_elements()) + elif isinstance(ele, TensorElement): + return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmety) + else: + return RestrictedElement(ele, restriction_domain) - W = FunctionSpace(V.mesh(), MixedElement(elements)) + W = FunctionSpace(V.mesh(), MixedElement([restrict(V.ufl_element(), d) for d in ("interior", "facet")])) assert W.dim() == V.dim(), "Dimensions of the original and decomposed spaces do not match" + mixed_operator = a(sum(TestFunctions(W)), sum(TrialFunctions(W)), coefficients={}) + mixed_bcs = tuple(bc.reconstruct(V=W[-1], g=0) for bc in bcs) + self.perm = None self.iperm = None indices = self.get_permutation(V, W) @@ -183,17 +195,6 @@ def initialize(self, pc): self.perm = PETSc.IS().createGeneral(indices, comm=V.comm) self.iperm = self.perm.invertPermutation() - mixed_operator = a(sum(TestFunctions(W)), sum(TrialFunctions(W)), coefficients={}) - mixed_bcs = tuple(bc.reconstruct(V=W[-1], g=0) for bc in problem.bcs) - - def _permute_nullspace(nsp): - if nsp is None or self.iperm is None: - return nsp - vecs = [vec.duplicate() for vec in nsp.getVecs()] - for vec in vecs: - vec.permute(self.iperm) - return PETSc.NullSpace().create(constant=nsp.constant, vectors=vecs, comm=nsp.comm) - if mat_type != "submatrix": self.mixed_op = allocate_matrix(mixed_operator, bcs=mixed_bcs, @@ -206,6 +207,14 @@ def _permute_nullspace(nsp): self._assemble_mixed_op() mixed_opmat = self.mixed_op.petscmat + def _permute_nullspace(nsp): + if nsp is None or self.iperm is None: + return nsp + vecs = [vec.duplicate() for vec in nsp.getVecs()] + for vec in vecs: + vec.permute(self.iperm) + return PETSc.NullSpace().create(constant=nsp.constant, vectors=vecs, comm=nsp.comm) + if False: # FIXME mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace())) diff --git a/tests/regression/test_facet_split.py b/tests/regression/test_facet_split.py new file mode 100644 index 0000000000..d633320be2 --- /dev/null +++ b/tests/regression/test_facet_split.py @@ -0,0 +1,38 @@ +import pytest +from firedrake import * + +@pytest.mark.parametrize(['quadrilateral'], + [(q,) + for q in [True, False]]) +def test_facet_split(quadrilateral): + parameters = { + "mat_type": "matfree", + "snes_type": "ksponly", + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.FacetSplitPC", + "facet": { + "mat_type": "aij", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_fact_type": "full", + "pc_fieldsplit_schur_precondition": "selfp", + "fieldsplit_0_pc_type": "jacobi", + "fieldsplit_1_pc_type": "lu", + }, + } + variant = "fdm" if quadrilateral else None + mesh = UnitSquareMesh(5, 5, quadrilateral=quadrilateral) + V = FunctionSpace(mesh, FiniteElement("Lagrange", degree=3, variant=variant)) + u = TrialFunction(V) + v = TestFunction(V) + uh = Function(V) + + a = inner(grad(u), grad(v))*dx + L = inner(v, Constant(0))*dx + x = SpatialCoordinate(mesh) + u_exact = 42 * x[1] + bcs = [DirichletBC(V, u_exact, "on_boundary")] + + solve(a == L, uh, bcs=bcs, solver_parameters=parameters) + assert sqrt(assemble(inner(uh - u_exact, uh - u_exact) * dx)) < 1E-10 From 7878f7ae2b828428b562e4a476c6407b0a9170c8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 16 Jan 2023 14:22:23 +0000 Subject: [PATCH 04/14] small changes --- firedrake/preconditioners/facet_split.py | 26 ++++++++++-------------- tests/regression/test_facet_split.py | 11 ++++++---- 2 files changed, 18 insertions(+), 19 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 491c0118a1..84282ee3c9 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -10,20 +10,18 @@ def split_dofs(elem): entity_dofs = elem.entity_dofs() ndim = elem.cell.get_spatial_dimension() - edofs = [[] for k in range(ndim+1)] - for key in entity_dofs: + edofs = [[], []] + for key in sorted(entity_dofs.keys()): vals = entity_dofs[key] edim = key try: edim = sum(edim) except TypeError: pass - split = numpy.arange(0, len(vals)+1, 2**(ndim-edim)) - for r in range(len(split)-1): - v = sum([vals[k] for k in range(*split[r:r+2])], []) - edofs[edim].extend(sorted(v)) + for k in vals: + edofs[edim < ndim].extend(sorted(vals[k])) - return numpy.array(edofs[-1], dtype=PETSc.IntType), numpy.array(sum(reversed(edofs[:-1]), []), dtype=PETSc.IntType) + return tuple(numpy.array(e, dtype=PETSc.IntType) for e in edofs) def restricted_dofs(celem, felem): @@ -161,7 +159,7 @@ def initialize(self, pc): options = PETSc.Options(options_prefix) mat_type = options.getString("mat_type", "submatrix") - if P.getType() == "python": + if P.getType() == "python" and False: ictx = P.getPythonContext() a = ictx.a bcs = tuple(ictx.row_bcs) @@ -208,18 +206,16 @@ def restrict(ele, restriction_domain): mixed_opmat = self.mixed_op.petscmat def _permute_nullspace(nsp): - if nsp is None or self.iperm is None: + if not (nsp.handle and self.iperm): return nsp vecs = [vec.duplicate() for vec in nsp.getVecs()] for vec in vecs: vec.permute(self.iperm) - return PETSc.NullSpace().create(constant=nsp.constant, vectors=vecs, comm=nsp.comm) + return PETSc.NullSpace().create(constant=nsp.hasConstant(), vectors=vecs, comm=nsp.getComm()) - if False: - # FIXME - mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace())) - mixed_opmat.setNearNullSpace(_permute_nullspace(P.getNearNullSpace())) - mixed_opmat.setTransposeNullSpace(_permute_nullspace(P.getTransposeNullSpace())) + mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace())) + mixed_opmat.setNearNullSpace(_permute_nullspace(P.getNearNullSpace())) + mixed_opmat.setTransposeNullSpace(_permute_nullspace(P.getTransposeNullSpace())) elif self.perm: self._permute_op = partial(PETSc.Mat().createSubMatrixVirtual, P, self.iperm, self.iperm) diff --git a/tests/regression/test_facet_split.py b/tests/regression/test_facet_split.py index d633320be2..75a78c3469 100644 --- a/tests/regression/test_facet_split.py +++ b/tests/regression/test_facet_split.py @@ -21,18 +21,21 @@ def test_facet_split(quadrilateral): "fieldsplit_1_pc_type": "lu", }, } + r = 2 variant = "fdm" if quadrilateral else None - mesh = UnitSquareMesh(5, 5, quadrilateral=quadrilateral) + mesh = UnitSquareMesh(2 ** r, 2 ** r, quadrilateral=quadrilateral) + V = FunctionSpace(mesh, FiniteElement("Lagrange", degree=3, variant=variant)) u = TrialFunction(V) v = TestFunction(V) uh = Function(V) - a = inner(grad(u), grad(v))*dx - L = inner(v, Constant(0))*dx + a = inner(grad(u), grad(v)) * dx + L = inner(Constant(0), v) * dx x = SpatialCoordinate(mesh) u_exact = 42 * x[1] - bcs = [DirichletBC(V, u_exact, "on_boundary")] + bcs = [DirichletBC(V, Constant(0), 3), + DirichletBC(V, Constant(42), 4)] solve(a == L, uh, bcs=bcs, solver_parameters=parameters) assert sqrt(assemble(inner(uh - u_exact, uh - u_exact) * dx)) < 1E-10 From 431efd47f1c7ae753cb5178c7294b26650bfa30e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 16 Jan 2023 14:39:28 +0000 Subject: [PATCH 05/14] lint --- tests/regression/test_facet_split.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/regression/test_facet_split.py b/tests/regression/test_facet_split.py index 75a78c3469..d8b2deaf66 100644 --- a/tests/regression/test_facet_split.py +++ b/tests/regression/test_facet_split.py @@ -1,6 +1,7 @@ import pytest from firedrake import * + @pytest.mark.parametrize(['quadrilateral'], [(q,) for q in [True, False]]) From 5ff48c8dda1c6f47b4d6756d8b6463b1551d0ead Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 16 Jan 2023 15:22:04 +0000 Subject: [PATCH 06/14] add documentation --- firedrake/preconditioners/facet_split.py | 216 ++++++++++++----------- tests/regression/test_facet_split.py | 63 ++++--- 2 files changed, 154 insertions(+), 125 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 84282ee3c9..aaee1e7b73 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -7,114 +7,17 @@ __all__ = ['FacetSplitPC'] -def split_dofs(elem): - entity_dofs = elem.entity_dofs() - ndim = elem.cell.get_spatial_dimension() - edofs = [[], []] - for key in sorted(entity_dofs.keys()): - vals = entity_dofs[key] - edim = key - try: - edim = sum(edim) - except TypeError: - pass - for k in vals: - edofs[edim < ndim].extend(sorted(vals[k])) - - return tuple(numpy.array(e, dtype=PETSc.IntType) for e in edofs) - - -def restricted_dofs(celem, felem): - """ - find which DOFs from felem are on celem - :arg celem: the restricted :class:`finat.FiniteElement` - :arg felem: the unrestricted :class:`finat.FiniteElement` - :returns: :class:`numpy.array` with indices of felem that correspond to celem - """ - csplit = split_dofs(celem) - fsplit = split_dofs(felem) - if len(csplit[0]) and len(csplit[1]): - csplit = [numpy.concatenate(csplit)] - fsplit = [numpy.concatenate(fsplit)] - - k = len(csplit[0]) == 0 - if len(csplit[k]) != len(fsplit[k]): - raise ValueError("Finite elements have different DOFs") - perm = numpy.empty_like(csplit[k]) - perm[csplit[k]] = numpy.arange(len(perm), dtype=perm.dtype) - return fsplit[k][perm] - - -def get_permutation_map(V, W): - from firedrake import Function - from pyop2 import op2, PermutedMap - - bsize = V.value_size - perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) - perm.fill(-1) - v = Function(V, dtype=PETSc.IntType, val=perm) - w = Function(W, dtype=PETSc.IntType) - - offset = 0 - for wdata, Wsub in zip(w.dat.data, W): - own = Wsub.dof_dset.set.size * bsize - wdata[:own] = numpy.arange(offset, offset+own, dtype=PETSc.IntType) - offset += own - - idofs = W[0].finat_element.space_dimension() * bsize - fdofs = W[1].finat_element.space_dimension() * bsize - - eperm = numpy.concatenate([restricted_dofs(Wsub.finat_element, V.finat_element) for Wsub in W]) - pmap = PermutedMap(V.cell_node_map(), eperm) - - kernel_code = f""" - void permutation(PetscInt *restrict x, - const PetscInt *restrict xi, - const PetscInt *restrict xf){{ - - for(PetscInt i=0; i<{idofs}; i++) x[i] = xi[i]; - for(PetscInt i=0; i<{fdofs}; i++) x[i+{idofs}] = xf[i]; - return; - }} - """ - kernel = op2.Kernel(kernel_code, "permutation", requires_zeroed_output_arguments=False) - op2.par_loop(kernel, v.cell_set, - v.dat(op2.WRITE, pmap), - w.dat[0](op2.READ, W[0].cell_node_map()), - w.dat[1](op2.READ, W[1].cell_node_map())) - - own = V.dof_dset.set.size * bsize - perm = V.dof_dset.lgmap.apply(perm, result=perm) - return perm[:own] - - -def get_permutation_project(V, W): - from firedrake import Function, split - ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() - start, end = ownership_ranges[V.comm.rank:V.comm.rank+2] - v = Function(V) - w = Function(W) - with w.dat.vec_wo as wvec: - wvec.setArray(numpy.linspace(0.0E0, 1.0E0, end-start, dtype=PETSc.RealType)) - w_expr = sum(split(w)) - try: - v.interpolate(w_expr) - except NotImplementedError: - rtol = 1.0 / max(numpy.diff(ownership_ranges))**2 - v.project(w_expr, solver_parameters={ - "mat_type": "matfree", - "ksp_type": "cg", - "ksp_atol": 0, - "ksp_rtol": rtol, - "pc_type": "jacobi", }) - return numpy.rint((end-start-1)*v.dat.data_ro.reshape((-1,))+start).astype(PETSc.IntType) - - class FacetSplitPC(PCBase): """ A preconditioner that splits a function into interior and facet DOFs. - Composability with fieldsplit can be done through the PETSc solver option: - -facet_pc_type fieldsplit + Internally this creates a PETSc PC object that can be controlled + by options using the extra options prefix ``facet_``. + + This allows for statically-condensed preconditioners to be applied to + linear systems involving the matrix applied to the full set of DOFs. Code + generated for the matrix-free operator evaluation in the space with full + DOFs will run faster than the one with interior-facet decoposition, since + the full element has a simpler structure. """ needs_python_pmat = False @@ -297,3 +200,106 @@ def destroy(self, pc): if hasattr(self, "perm"): if self.perm: self.perm.destroy() + + +def split_dofs(elem): + entity_dofs = elem.entity_dofs() + ndim = elem.cell.get_spatial_dimension() + edofs = [[], []] + for key in sorted(entity_dofs.keys()): + vals = entity_dofs[key] + edim = key + try: + edim = sum(edim) + except TypeError: + pass + for k in vals: + edofs[edim < ndim].extend(sorted(vals[k])) + + return tuple(numpy.array(e, dtype=PETSc.IntType) for e in edofs) + + +def restricted_dofs(celem, felem): + """ + find which DOFs from felem are on celem + :arg celem: the restricted :class:`finat.FiniteElement` + :arg felem: the unrestricted :class:`finat.FiniteElement` + :returns: :class:`numpy.array` with indices of felem that correspond to celem + """ + csplit = split_dofs(celem) + fsplit = split_dofs(felem) + if len(csplit[0]) and len(csplit[1]): + csplit = [numpy.concatenate(csplit)] + fsplit = [numpy.concatenate(fsplit)] + + k = len(csplit[0]) == 0 + if len(csplit[k]) != len(fsplit[k]): + raise ValueError("Finite elements have different DOFs") + perm = numpy.empty_like(csplit[k]) + perm[csplit[k]] = numpy.arange(len(perm), dtype=perm.dtype) + return fsplit[k][perm] + + +def get_permutation_map(V, W): + from firedrake import Function + from pyop2 import op2, PermutedMap + + bsize = V.value_size + idofs = W[0].finat_element.space_dimension() * bsize + fdofs = W[1].finat_element.space_dimension() * bsize + + perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) + perm.fill(-1) + v = Function(V, dtype=PETSc.IntType, val=perm) + w = Function(W, dtype=PETSc.IntType) + + offset = 0 + for wdata, Wsub in zip(w.dat.data, W): + own = Wsub.dof_dset.set.size * bsize + wdata[:own] = numpy.arange(offset, offset+own, dtype=PETSc.IntType) + offset += own + + eperm = numpy.concatenate([restricted_dofs(Wsub.finat_element, V.finat_element) for Wsub in W]) + pmap = PermutedMap(V.cell_node_map(), eperm) + + kernel_code = f""" + void permutation(PetscInt *restrict x, + const PetscInt *restrict xi, + const PetscInt *restrict xf){{ + + for(PetscInt i=0; i<{idofs}; i++) x[i] = xi[i]; + for(PetscInt i=0; i<{fdofs}; i++) x[i+{idofs}] = xf[i]; + return; + }} + """ + kernel = op2.Kernel(kernel_code, "permutation", requires_zeroed_output_arguments=False) + op2.par_loop(kernel, v.cell_set, + v.dat(op2.WRITE, pmap), + w.dat[0](op2.READ, W[0].cell_node_map()), + w.dat[1](op2.READ, W[1].cell_node_map())) + + own = V.dof_dset.set.size * bsize + perm = V.dof_dset.lgmap.apply(perm, result=perm) + return perm[:own] + + +def get_permutation_project(V, W): + from firedrake import Function, split + ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() + start, end = ownership_ranges[V.comm.rank:V.comm.rank+2] + v = Function(V) + w = Function(W) + with w.dat.vec_wo as wvec: + wvec.setArray(numpy.linspace(0.0E0, 1.0E0, end-start, dtype=PETSc.RealType)) + w_expr = sum(split(w)) + try: + v.interpolate(w_expr) + except NotImplementedError: + rtol = 1.0 / max(numpy.diff(ownership_ranges))**2 + v.project(w_expr, solver_parameters={ + "mat_type": "matfree", + "ksp_type": "cg", + "ksp_atol": 0, + "ksp_rtol": rtol, + "pc_type": "jacobi", }) + return numpy.rint((end-start-1)*v.dat.data_ro.reshape((-1,))+start).astype(PETSc.IntType) diff --git a/tests/regression/test_facet_split.py b/tests/regression/test_facet_split.py index d8b2deaf66..3634f07a85 100644 --- a/tests/regression/test_facet_split.py +++ b/tests/regression/test_facet_split.py @@ -2,26 +2,49 @@ from firedrake import * -@pytest.mark.parametrize(['quadrilateral'], - [(q,) - for q in [True, False]]) -def test_facet_split(quadrilateral): - parameters = { - "mat_type": "matfree", - "snes_type": "ksponly", - "ksp_type": "preonly", - "pc_type": "python", - "pc_python_type": "firedrake.FacetSplitPC", - "facet": { - "mat_type": "aij", - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "schur", - "pc_fieldsplit_schur_fact_type": "full", - "pc_fieldsplit_schur_precondition": "selfp", - "fieldsplit_0_pc_type": "jacobi", - "fieldsplit_1_pc_type": "lu", - }, - } +@pytest.mark.parametrize(['quadrilateral', 'ptype'], + [(q, p) + for q in [True, False] + for p in ["lu", "jacobi"]]) +def test_facet_split(quadrilateral, ptype): + if ptype == "lu": + parameters = { + "mat_type": "matfree", + "snes_type": "ksponly", + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.FacetSplitPC", + "facet": { + "mat_type": "aij", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_fact_type": "full", + "pc_fieldsplit_schur_precondition": "selfp", + "fieldsplit_0_pc_type": "jacobi", + "fieldsplit_1_pc_type": "lu", + }, + } + elif ptype == "jacobi": + parameters = { + "mat_type": "matfree", + "snes_type": "ksponly", + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.FacetSplitPC", + "facet": { + "mat_type": "submatrix", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_fact_type": "full", + "pc_fieldsplit_schur_precondition": "a11", + "fieldsplit_0_pc_type": "jacobi", + "fieldsplit_1_pc_type": "jacobi", + "fieldsplit_1_ksp_type": "cg", + "fieldsplit_1_ksp_rtol": 1E-12, + "fieldsplit_1_ksp_atol": 1E-10, + }, + } + r = 2 variant = "fdm" if quadrilateral else None mesh = UnitSquareMesh(2 ** r, 2 ** r, quadrilateral=quadrilateral) From 7ad2ef7404d2dbc2c7441724135f1482a2da8644 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 16 Jan 2023 17:28:01 +0000 Subject: [PATCH 07/14] clean up --- firedrake/preconditioners/facet_split.py | 28 ++++++------------------ tests/regression/test_facet_split.py | 2 -- 2 files changed, 7 insertions(+), 23 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index aaee1e7b73..c5904e474c 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -37,39 +37,28 @@ def get_permutation(self, V, W): return self._permutation_cache[key] def initialize(self, pc): - from ufl import RestrictedElement, MixedElement, TensorElement, VectorElement from firedrake import FunctionSpace, TestFunctions, TrialFunctions from firedrake.assemble import allocate_matrix, TwoFormAssembler - from firedrake.solving_utils import _SNESContext from functools import partial _, P = pc.getOperators() appctx = self.get_appctx(pc) fcp = appctx.get("form_compiler_parameters") - if pc.getType() != "python": - raise ValueError("Expecting PC type python") - - ctx = dmhooks.get_appctx(pc.getDM()) - if ctx is None: - raise ValueError("No context found.") - if not isinstance(ctx, _SNESContext): - raise ValueError("Don't know how to get form from %r", ctx) - prefix = pc.getOptionsPrefix() options_prefix = prefix + self._prefix options = PETSc.Options(options_prefix) mat_type = options.getString("mat_type", "submatrix") - if P.getType() == "python" and False: - ictx = P.getPythonContext() - a = ictx.a - bcs = tuple(ictx.row_bcs) + if P.getType() == "python": + ctx = P.getPythonContext() + a = ctx.a + bcs = tuple(ctx.row_bcs) else: - problem = ctx._problem - a = problem.Jp or problem.J - bcs = tuple(problem.bcs) + ctx = dmhooks.get_appctx(pc.getDM()) + a = ctx.Jp or ctx.J + bcs = tuple(ctx._problem.bcs) V = a.arguments()[-1].function_space() assert len(V) == 1, "Interior-facet decomposition of mixed elements is not supported" @@ -119,7 +108,6 @@ def _permute_nullspace(nsp): mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace())) mixed_opmat.setNearNullSpace(_permute_nullspace(P.getNearNullSpace())) mixed_opmat.setTransposeNullSpace(_permute_nullspace(P.getTransposeNullSpace())) - elif self.perm: self._permute_op = partial(PETSc.Mat().createSubMatrixVirtual, P, self.iperm, self.iperm) mixed_opmat = self._permute_op() @@ -215,7 +203,6 @@ def split_dofs(elem): pass for k in vals: edofs[edim < ndim].extend(sorted(vals[k])) - return tuple(numpy.array(e, dtype=PETSc.IntType) for e in edofs) @@ -266,7 +253,6 @@ def get_permutation_map(V, W): void permutation(PetscInt *restrict x, const PetscInt *restrict xi, const PetscInt *restrict xf){{ - for(PetscInt i=0; i<{idofs}; i++) x[i] = xi[i]; for(PetscInt i=0; i<{fdofs}; i++) x[i+{idofs}] = xf[i]; return; diff --git a/tests/regression/test_facet_split.py b/tests/regression/test_facet_split.py index 3634f07a85..bcb9e85e2b 100644 --- a/tests/regression/test_facet_split.py +++ b/tests/regression/test_facet_split.py @@ -10,7 +10,6 @@ def test_facet_split(quadrilateral, ptype): if ptype == "lu": parameters = { "mat_type": "matfree", - "snes_type": "ksponly", "ksp_type": "preonly", "pc_type": "python", "pc_python_type": "firedrake.FacetSplitPC", @@ -27,7 +26,6 @@ def test_facet_split(quadrilateral, ptype): elif ptype == "jacobi": parameters = { "mat_type": "matfree", - "snes_type": "ksponly", "ksp_type": "preonly", "pc_type": "python", "pc_python_type": "firedrake.FacetSplitPC", From 6ea8ffd641f72fe18bfcc71961c13e9c69350493 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Jan 2023 13:47:06 +0000 Subject: [PATCH 08/14] add comment --- firedrake/preconditioners/facet_split.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index c5904e474c..b69bc2ca78 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -63,7 +63,6 @@ def initialize(self, pc): V = a.arguments()[-1].function_space() assert len(V) == 1, "Interior-facet decomposition of mixed elements is not supported" - # W = V[interior] * V[facet] def restrict(ele, restriction_domain): if isinstance(ele, VectorElement): return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_elements()) @@ -72,6 +71,7 @@ def restrict(ele, restriction_domain): else: return RestrictedElement(ele, restriction_domain) + # W = V[interior] * V[facet] W = FunctionSpace(V.mesh(), MixedElement([restrict(V.ufl_element(), d) for d in ("interior", "facet")])) assert W.dim() == V.dim(), "Dimensions of the original and decomposed spaces do not match" @@ -191,6 +191,8 @@ def destroy(self, pc): def split_dofs(elem): + """ Split DOFs into interior and facet DOF, where facets are sorted by entity. + """ entity_dofs = elem.entity_dofs() ndim = elem.cell.get_spatial_dimension() edofs = [[], []] @@ -207,8 +209,7 @@ def split_dofs(elem): def restricted_dofs(celem, felem): - """ - find which DOFs from felem are on celem + """ Find which DOFs from felem are on celem :arg celem: the restricted :class:`finat.FiniteElement` :arg felem: the unrestricted :class:`finat.FiniteElement` :returns: :class:`numpy.array` with indices of felem that correspond to celem @@ -270,6 +271,8 @@ def get_permutation_map(V, W): def get_permutation_project(V, W): + """ Alternative projection-based method to obtain DOF permutation + """ from firedrake import Function, split ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() start, end = ownership_ranges[V.comm.rank:V.comm.rank+2] From 3d5256c75818ae74644cf260eb1aa76b095d7380 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Jan 2023 15:19:05 +0000 Subject: [PATCH 09/14] clean up --- firedrake/preconditioners/facet_split.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index b69bc2ca78..8a143e3492 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -232,10 +232,6 @@ def get_permutation_map(V, W): from firedrake import Function from pyop2 import op2, PermutedMap - bsize = V.value_size - idofs = W[0].finat_element.space_dimension() * bsize - fdofs = W[1].finat_element.space_dimension() * bsize - perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) perm.fill(-1) v = Function(V, dtype=PETSc.IntType, val=perm) @@ -243,10 +239,11 @@ def get_permutation_map(V, W): offset = 0 for wdata, Wsub in zip(w.dat.data, W): - own = Wsub.dof_dset.set.size * bsize + own = Wsub.dof_dset.set.size * Wsub.value_size wdata[:own] = numpy.arange(offset, offset+own, dtype=PETSc.IntType) offset += own + sizes = [Wsub.finat_element.space_dimension() * Wsub.value_size for Wsub in W] eperm = numpy.concatenate([restricted_dofs(Wsub.finat_element, V.finat_element) for Wsub in W]) pmap = PermutedMap(V.cell_node_map(), eperm) @@ -254,8 +251,8 @@ def get_permutation_map(V, W): void permutation(PetscInt *restrict x, const PetscInt *restrict xi, const PetscInt *restrict xf){{ - for(PetscInt i=0; i<{idofs}; i++) x[i] = xi[i]; - for(PetscInt i=0; i<{fdofs}; i++) x[i+{idofs}] = xf[i]; + for(PetscInt i=0; i<{sizes[0]}; i++) x[i] = xi[i]; + for(PetscInt i=0; i<{sizes[1]}; i++) x[i+{sizes[0]}] = xf[i]; return; }} """ @@ -265,7 +262,7 @@ def get_permutation_map(V, W): w.dat[0](op2.READ, W[0].cell_node_map()), w.dat[1](op2.READ, W[1].cell_node_map())) - own = V.dof_dset.set.size * bsize + own = V.dof_dset.set.size * V.value_size perm = V.dof_dset.lgmap.apply(perm, result=perm) return perm[:own] From d4d9157eff29613a44badc7a6067087a43430b36 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Jan 2023 16:20:10 +0000 Subject: [PATCH 10/14] fix typo --- firedrake/preconditioners/facet_split.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 8a143e3492..e56d0df6fe 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -65,7 +65,7 @@ def initialize(self, pc): def restrict(ele, restriction_domain): if isinstance(ele, VectorElement): - return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_elements()) + return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_sub_elements()) elif isinstance(ele, TensorElement): return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmety) else: From 6f9884423a76866baa93708dfffb702cc9642656 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 19 Jan 2023 17:07:58 +0000 Subject: [PATCH 11/14] fix permutation for vector spaces --- firedrake/preconditioners/facet_split.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index e56d0df6fe..4c6fa690ed 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -235,13 +235,14 @@ def get_permutation_map(V, W): perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) perm.fill(-1) v = Function(V, dtype=PETSc.IntType, val=perm) - w = Function(W, dtype=PETSc.IntType) offset = 0 - for wdata, Wsub in zip(w.dat.data, W): - own = Wsub.dof_dset.set.size * Wsub.value_size - wdata[:own] = numpy.arange(offset, offset+own, dtype=PETSc.IntType) - offset += own + wdats = [] + for Wsub in W: + data = numpy.arange(offset, offset + Wsub.dof_count, dtype=PETSc.IntType) + wsub = Function(Wsub, dtype=PETSc.IntType, val=data) + wdats.append(wsub.dat) + offset += Wsub.dof_dset.layout_vec.sizes[0] sizes = [Wsub.finat_element.space_dimension() * Wsub.value_size for Wsub in W] eperm = numpy.concatenate([restricted_dofs(Wsub.finat_element, V.finat_element) for Wsub in W]) @@ -259,10 +260,11 @@ def get_permutation_map(V, W): kernel = op2.Kernel(kernel_code, "permutation", requires_zeroed_output_arguments=False) op2.par_loop(kernel, v.cell_set, v.dat(op2.WRITE, pmap), - w.dat[0](op2.READ, W[0].cell_node_map()), - w.dat[1](op2.READ, W[1].cell_node_map())) + wdats[0](op2.READ, W[0].cell_node_map()), + wdats[1](op2.READ, W[1].cell_node_map())) - own = V.dof_dset.set.size * V.value_size + own = V.dof_dset.layout_vec.sizes[0] + perm = perm.reshape((-1, )) perm = V.dof_dset.lgmap.apply(perm, result=perm) return perm[:own] From c15af29ba34ad37efcdeb18dcce95a69fa9c571d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 25 Jan 2023 14:26:59 +0000 Subject: [PATCH 12/14] get rid of Functions --- firedrake/preconditioners/facet_split.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 4c6fa690ed..be39f80aac 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -229,19 +229,17 @@ def restricted_dofs(celem, felem): def get_permutation_map(V, W): - from firedrake import Function from pyop2 import op2, PermutedMap perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) perm.fill(-1) - v = Function(V, dtype=PETSc.IntType, val=perm) + vdat = V.make_dat(val=perm) offset = 0 wdats = [] for Wsub in W: - data = numpy.arange(offset, offset + Wsub.dof_count, dtype=PETSc.IntType) - wsub = Function(Wsub, dtype=PETSc.IntType, val=data) - wdats.append(wsub.dat) + val = numpy.arange(offset, offset + Wsub.dof_count, dtype=PETSc.IntType) + wdats.append(Wsub.make_dat(val=val)) offset += Wsub.dof_dset.layout_vec.sizes[0] sizes = [Wsub.finat_element.space_dimension() * Wsub.value_size for Wsub in W] @@ -258,8 +256,8 @@ def get_permutation_map(V, W): }} """ kernel = op2.Kernel(kernel_code, "permutation", requires_zeroed_output_arguments=False) - op2.par_loop(kernel, v.cell_set, - v.dat(op2.WRITE, pmap), + op2.par_loop(kernel, V.mesh().cell_set, + vdat(op2.WRITE, pmap), wdats[0](op2.READ, W[0].cell_node_map()), wdats[1](op2.READ, W[1].cell_node_map())) From a32319e6eae163374f8c8b42ae3a75150669f8f1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 25 Jan 2023 14:51:33 +0000 Subject: [PATCH 13/14] V.comm -> V._comm, move some imports to the top --- firedrake/preconditioners/facet_split.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index be39f80aac..824c9f72ac 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -1,9 +1,11 @@ +from functools import partial +from mpi4py import MPI +from pyop2 import op2, PermutedMap from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase import firedrake.dmhooks as dmhooks import numpy - __all__ = ['FacetSplitPC'] @@ -26,11 +28,10 @@ class FacetSplitPC(PCBase): _permutation_cache = {} def get_permutation(self, V, W): - from mpi4py import MPI key = (V, W) if key not in self._permutation_cache: indices = get_permutation_map(V, W) - if V.comm.allreduce(numpy.all(indices[:-1] <= indices[1:]), MPI.PROD): + if V._comm.allreduce(numpy.all(indices[:-1] <= indices[1:]), MPI.PROD): self._permutation_cache[key] = None else: self._permutation_cache[key] = indices @@ -40,7 +41,6 @@ def initialize(self, pc): from ufl import RestrictedElement, MixedElement, TensorElement, VectorElement from firedrake import FunctionSpace, TestFunctions, TrialFunctions from firedrake.assemble import allocate_matrix, TwoFormAssembler - from functools import partial _, P = pc.getOperators() appctx = self.get_appctx(pc) @@ -82,7 +82,7 @@ def restrict(ele, restriction_domain): self.iperm = None indices = self.get_permutation(V, W) if indices is not None: - self.perm = PETSc.IS().createGeneral(indices, comm=V.comm) + self.perm = PETSc.IS().createGeneral(indices, comm=V._comm) self.iperm = self.perm.invertPermutation() if mat_type != "submatrix": @@ -203,7 +203,7 @@ def split_dofs(elem): edim = sum(edim) except TypeError: pass - for k in vals: + for k in sorted(vals.keys()): edofs[edim < ndim].extend(sorted(vals[k])) return tuple(numpy.array(e, dtype=PETSc.IntType) for e in edofs) @@ -229,8 +229,6 @@ def restricted_dofs(celem, felem): def get_permutation_map(V, W): - from pyop2 import op2, PermutedMap - perm = numpy.empty((V.dof_count, ), dtype=PETSc.IntType) perm.fill(-1) vdat = V.make_dat(val=perm) @@ -270,9 +268,10 @@ def get_permutation_map(V, W): def get_permutation_project(V, W): """ Alternative projection-based method to obtain DOF permutation """ - from firedrake import Function, split + from firedrake import Function + from ufl import split ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() - start, end = ownership_ranges[V.comm.rank:V.comm.rank+2] + start, end = ownership_ranges[V._comm.rank:V._comm.rank+2] v = Function(V) w = Function(W) with w.dat.vec_wo as wvec: From 2c64af8b7cd7cbcc86a3414e6079199272172c48 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 25 Jan 2023 16:46:33 +0000 Subject: [PATCH 14/14] remove dead code --- firedrake/preconditioners/facet_split.py | 25 ------------------------ 1 file changed, 25 deletions(-) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 824c9f72ac..a1d134f206 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -263,28 +263,3 @@ def get_permutation_map(V, W): perm = perm.reshape((-1, )) perm = V.dof_dset.lgmap.apply(perm, result=perm) return perm[:own] - - -def get_permutation_project(V, W): - """ Alternative projection-based method to obtain DOF permutation - """ - from firedrake import Function - from ufl import split - ownership_ranges = V.dof_dset.layout_vec.getOwnershipRanges() - start, end = ownership_ranges[V._comm.rank:V._comm.rank+2] - v = Function(V) - w = Function(W) - with w.dat.vec_wo as wvec: - wvec.setArray(numpy.linspace(0.0E0, 1.0E0, end-start, dtype=PETSc.RealType)) - w_expr = sum(split(w)) - try: - v.interpolate(w_expr) - except NotImplementedError: - rtol = 1.0 / max(numpy.diff(ownership_ranges))**2 - v.project(w_expr, solver_parameters={ - "mat_type": "matfree", - "ksp_type": "cg", - "ksp_atol": 0, - "ksp_rtol": rtol, - "pc_type": "jacobi", }) - return numpy.rint((end-start-1)*v.dat.data_ro.reshape((-1,))+start).astype(PETSc.IntType)