diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 6cd6f3d5d7..0d16b0f7cb 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -319,6 +319,8 @@ def reconstruct(self, field=None, V=None, g=None, sub_domain=None, use_split=Fal V = V.sub(index) if g is None: g = self._original_arg + if isinstance(g, firedrake.Function) and g.function_space() != V: + g = firedrake.Function(V).interpolate(g) if sub_domain is None: sub_domain = self.sub_domain if field is not None: @@ -739,11 +741,11 @@ def restricted_function_space(V, ids): return V assert len(ids) == len(V) - spaces = [Vsub if len(boundary_set) == 0 else - firedrake.RestrictedFunctionSpace(Vsub, boundary_set=boundary_set) - for Vsub, boundary_set in zip(V, ids)] + spaces = [V_ if len(boundary_set) == 0 else + firedrake.RestrictedFunctionSpace(V_, boundary_set=boundary_set, name=V_.name) + for V_, boundary_set in zip(V, ids)] if len(spaces) == 1: return spaces[0] else: - return firedrake.MixedFunctionSpace(spaces) + return firedrake.MixedFunctionSpace(spaces, name=V.name) diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index c1cca65b56..07b334b111 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -53,11 +53,13 @@ def get_function_space(dm): :raises RuntimeError: if no function space was found. """ info = dm.getAttr("__fs_info__") - meshref, element, indices, (name, names) = info + meshref, element, indices, (name, names), boundary_sets = info mesh = meshref() if mesh is None: raise RuntimeError("Somehow your mesh was collected, this should never happen") V = firedrake.FunctionSpace(mesh, element, name=name) + if any(boundary_sets): + V = firedrake.bcs.restricted_function_space(V, boundary_sets) if len(V) > 1: for V_, name in zip(V, names): V_.topological.name = name @@ -93,8 +95,8 @@ def set_function_space(dm, V): if len(V) > 1: names = tuple(V_.name for V_ in V) element = V.ufl_element() - - info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names)) + boundary_sets = tuple(V_.boundary_set for V_ in V) + info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names), boundary_sets) dm.setAttr("__fs_info__", info) @@ -457,7 +459,7 @@ def refine(dm, comm): if hasattr(V, "_fine"): fdm = V._fine.dm else: - V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element()) + V._fine = V.reconstruct(mesh=hierarchy[level + 1]) fdm = V._fine.dm V._fine._coarse = V return fdm diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index cec9dd860c..60d518f884 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -378,10 +378,13 @@ def make_function_space(cls, mesh, element, name=None): new = cls.create(new, mesh) return new - def reconstruct(self, mesh=None, name=None, **kwargs): + def reconstruct(self, mesh=None, element=None, boundary_set=None, name=None, **kwargs): r"""Reconstruct this :class:`.WithGeometryBase` . :kwarg mesh: the new :func:`~.Mesh` (defaults to same mesh) + :kwarg element: the new :class:`finat.ufl.FiniteElement` (defaults to same element) + :kwarg boundary_set: boundary subdomain labels defining a new + :func:`~.RestrictedFunctionSpace` (defaults to same boundary_set) :kwarg name: the new name (defaults to None) :returns: the new function space of the same class as ``self``. @@ -389,6 +392,9 @@ def reconstruct(self, mesh=None, name=None, **kwargs): For details see :meth:`finat.ufl.finiteelement.FiniteElement.reconstruct`. """ V_parent = self + if boundary_set is None: + boundary_set = V_parent.boundary_set + # Deal with ProxyFunctionSpace indices = [] while True: @@ -403,8 +409,9 @@ def reconstruct(self, mesh=None, name=None, **kwargs): if mesh is None: mesh = V_parent.mesh() + if element is None: + element = V_parent.ufl_element() - element = V_parent.ufl_element() cell = mesh.topology.ufl_cell() if len(kwargs) > 0 or element.cell != cell: element = element.reconstruct(cell=cell, **kwargs) @@ -412,6 +419,9 @@ def reconstruct(self, mesh=None, name=None, **kwargs): V = type(self).make_function_space(mesh, element, name=name) for i in reversed(indices): V = V.sub(i) + + if boundary_set: + V = RestrictedFunctionSpace(V, boundary_set=boundary_set, name=V.name) return V @@ -876,6 +886,14 @@ class RestrictedFunctionSpace(FunctionSpace): If using this class to solve or similar, a list of DirichletBCs will still need to be specified on this space and passed into the function. """ + def __new__(cls, function_space, boundary_set=frozenset(), name=None): + mesh = function_space.mesh() + if mesh is not mesh.topology: + V = RestrictedFunctionSpace(function_space.topological, + boundary_set=boundary_set, name=name) + return type(function_space).create(V, mesh) + return FunctionSpace.__new__(cls) + def __init__(self, function_space, boundary_set=frozenset(), name=None): label = "" boundary_set_ = [] @@ -901,8 +919,7 @@ def __init__(self, function_space, boundary_set=frozenset(), name=None): function_space.ufl_element(), label=self._label) self.function_space = function_space - self.name = name or (function_space.name or "Restricted" + "_" - + "_".join(sorted(map(str, self.boundary_set)))) + self.name = name or function_space.name def set_shared_data(self): sdata = get_shared_data(self._mesh, self.ufl_element(), self.boundary_set) @@ -1200,7 +1217,7 @@ class ProxyFunctionSpace(FunctionSpace): """ def __new__(cls, mesh, element, name=None): topology = mesh.topology - self = super(ProxyFunctionSpace, cls).__new__(cls) + self = FunctionSpace.__new__(cls) if mesh is not topology: return WithGeometry.create(self, mesh) else: @@ -1255,7 +1272,7 @@ class ProxyRestrictedFunctionSpace(RestrictedFunctionSpace): """ def __new__(cls, function_space, boundary_set=frozenset(), name=None): topology = function_space._mesh.topology - self = super(ProxyRestrictedFunctionSpace, cls).__new__(cls) + self = FunctionSpace.__new__(cls) if function_space._mesh is not topology: return WithGeometry.create(self, function_space._mesh) else: diff --git a/firedrake/mg/embedded.py b/firedrake/mg/embedded.py index 551c09a7e2..17b8af37eb 100644 --- a/firedrake/mg/embedded.py +++ b/firedrake/mg/embedded.py @@ -90,6 +90,9 @@ def cache(self, V): except KeyError: return self.caches.setdefault(key, TransferManager.Cache(*key)) + def cache_key(self, V): + return (V.dim(), tuple(V.boundary_set)) + def V_dof_weights(self, V): """Dof weights for averaging projection. @@ -97,7 +100,7 @@ def V_dof_weights(self, V): :returns: A PETSc Vec. """ cache = self.cache(V) - key = V.dim() + key = self.cache_key(V) try: return cache._V_dof_weights[key] except KeyError: @@ -122,7 +125,7 @@ def V_DG_mass(self, V, DG): :returns: A PETSc Mat mapping from V -> DG """ cache = self.cache(V) - key = V.dim() + key = self.cache_key(V) try: return cache._V_DG_mass[key] except KeyError: @@ -153,7 +156,7 @@ def V_approx_inv_mass(self, V, DG): :returns: A PETSc Mat mapping from V -> DG. """ cache = self.cache(V) - key = V.dim() + key = self.cache_key(V) try: return cache._V_approx_inv_mass[key] except KeyError: @@ -171,7 +174,7 @@ def V_inv_mass_ksp(self, V): :returns: A PETSc KSP for inverting (V, V). """ cache = self.cache(V) - key = V.dim() + key = self.cache_key(V) try: return cache._V_inv_mass_ksp[key] except KeyError: @@ -193,7 +196,7 @@ def DG_work(self, V): """ needs_dual = ufl.duals.is_dual(V) cache = self.cache(V) - key = (V.dim(), needs_dual) + key = self.cache_key(V) + (needs_dual,) try: return cache._DG_work[key] except KeyError: @@ -210,7 +213,7 @@ def work_vec(self, V): :returns: A PETSc Vec for V. """ cache = self.cache(V) - key = V.dim() + key = self.cache_key(V) try: return cache._work_vec[key] except KeyError: diff --git a/firedrake/mg/interface.py b/firedrake/mg/interface.py index f962253db3..22faeb3e8d 100644 --- a/firedrake/mg/interface.py +++ b/firedrake/mg/interface.py @@ -58,7 +58,6 @@ def prolong(coarse, fine): repeat = (fine_level - coarse_level)*refinements_per_level next_level = coarse_level * refinements_per_level - element = Vc.ufl_element() meshes = hierarchy._meshes for j in range(repeat): next_level += 1 @@ -66,10 +65,10 @@ def prolong(coarse, fine): next = fine Vf = fine.function_space() else: - Vf = firedrake.FunctionSpace(meshes[next_level], element) + Vf = Vc.reconstruct(mesh=meshes[next_level]) next = firedrake.Function(Vf) - coarse_coords = Vc.mesh().coordinates + coarse_coords = get_coordinates(Vc) fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc) fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space()) kernel = kernels.prolong_kernel(coarse) @@ -119,7 +118,6 @@ def restrict(fine_dual, coarse_dual): repeat = (fine_level - coarse_level)*refinements_per_level next_level = fine_level * refinements_per_level - element = Vc.ufl_element() meshes = hierarchy._meshes for j in range(repeat): @@ -128,15 +126,15 @@ def restrict(fine_dual, coarse_dual): coarse_dual.dat.zero() next = coarse_dual else: - Vc = firedrake.FunctionSpace(meshes[next_level], element) - next = firedrake.Cofunction(Vc.dual()) + Vc = Vf.reconstruct(mesh=meshes[next_level]) + next = firedrake.Cofunction(Vc) Vc = next.function_space() # XXX: Should be able to figure out locations by pushing forward # reference cell node locations to physical space. # x = \sum_i c_i \phi_i(x_hat) - node_locations = utils.physical_node_locations(Vf) + node_locations = utils.physical_node_locations(Vf.dual()) - coarse_coords = Vc.mesh().coordinates + coarse_coords = get_coordinates(Vc.dual()) fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc) fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space()) # Have to do this, because the node set core size is not right for @@ -195,7 +193,6 @@ def inject(fine, coarse): repeat = (fine_level - coarse_level)*refinements_per_level next_level = fine_level * refinements_per_level - element = Vc.ufl_element() meshes = hierarchy._meshes for j in range(repeat): @@ -205,12 +202,12 @@ def inject(fine, coarse): next = coarse Vc = next.function_space() else: - Vc = firedrake.FunctionSpace(meshes[next_level], element) + Vc = Vf.reconstruct(mesh=meshes[next_level]) next = firedrake.Function(Vc) if not dg: node_locations = utils.physical_node_locations(Vc) - fine_coords = Vf.mesh().coordinates + fine_coords = get_coordinates(Vf) coarse_node_to_fine_nodes = utils.coarse_node_to_fine_node_map(Vc, Vf) coarse_node_to_fine_coords = utils.coarse_node_to_fine_node_map(Vc, fine_coords.function_space()) @@ -242,3 +239,11 @@ def inject(fine, coarse): fine = next Vf = Vc return coarse + + +def get_coordinates(V): + coords = V.mesh().coordinates + if V.boundary_set: + W = V.reconstruct(element=coords.function_space().ufl_element()) + coords = firedrake.Function(W).interpolate(coords) + return coords diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 9e92159d25..6786ba1f43 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -186,7 +186,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse): if isinstance(g, firedrake.Function) and hasattr(g, "_child"): manager.inject(g, g._child) - V = problem.u.function_space() + V = problem.u_restrict.function_space() if not hasattr(V, "_coarse"): # The hook is persistent and cumulative, but also problem-independent. # Therefore, we are only adding it once. @@ -201,7 +201,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse): for c in coefficients: coefficient_mapping[c] = self(c, self, coefficient_mapping=coefficient_mapping) - u = coefficient_mapping[problem.u] + u = coefficient_mapping[problem.u_restrict] bcs = [self(bc, self) for bc in problem.bcs] J = self(problem.J, self, coefficient_mapping=coefficient_mapping) @@ -277,7 +277,7 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): if isinstance(val, (firedrake.Function, firedrake.Cofunction)): V = val.function_space() coarseneddm = V.dm - parentdm = get_parent(context._problem.u.function_space().dm) + parentdm = get_parent(context._problem.u_restrict.function_space().dm) # Now attach the hook to the parent DM if get_appctx(coarseneddm) is None: @@ -369,8 +369,8 @@ def create_interpolation(dmc, dmf): manager = get_transfer_manager(dmf) - V_c = cctx._problem.u.function_space() - V_f = fctx._problem.u.function_space() + V_c = cctx._problem.u_restrict.function_space() + V_f = fctx._problem.u_restrict.function_space() row_size = V_f.dof_dset.layout_vec.getSizes() col_size = V_c.dof_dset.layout_vec.getSizes() @@ -395,8 +395,8 @@ def create_injection(dmc, dmf): manager = get_transfer_manager(dmf) - V_c = cctx._problem.u.function_space() - V_f = fctx._problem.u.function_space() + V_c = cctx._problem.u_restrict.function_space() + V_f = fctx._problem.u_restrict.function_space() row_size = V_f.dof_dset.layout_vec.getSizes() col_size = V_c.dof_dset.layout_vec.getSizes() diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index 80e5e820f6..025e5e462f 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -11,7 +11,7 @@ def fine_node_to_coarse_node_map(Vf, Vc): if len(Vf) > 1: assert len(Vf) == len(Vc) - return op2.MixedMap(fine_node_to_coarse_node_map(f, c) for f, c in zip(Vf, Vc)) + return op2.MixedMap(map(fine_node_to_coarse_node_map, Vf, Vc)) mesh = Vf.mesh() assert hasattr(mesh, "_shared_data_cache") hierarchyf, levelf = get_level(Vf.mesh()) @@ -25,10 +25,7 @@ def fine_node_to_coarse_node_map(Vf, Vc): if levelc + increment != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) - key = (entity_dofs_key(Vc.finat_element.entity_dofs()) - + entity_dofs_key(Vf.finat_element.entity_dofs()) - + (levelc, levelf)) - + key = _cache_key(Vc, Vf) cache = mesh._shared_data_cache["hierarchy_fine_node_to_coarse_node_map"] try: return cache[key] @@ -49,7 +46,7 @@ def fine_node_to_coarse_node_map(Vf, Vc): def coarse_node_to_fine_node_map(Vc, Vf): if len(Vf) > 1: assert len(Vf) == len(Vc) - return op2.MixedMap(coarse_node_to_fine_node_map(f, c) for f, c in zip(Vf, Vc)) + return op2.MixedMap(map(coarse_node_to_fine_node_map, Vf, Vc)) mesh = Vc.mesh() assert hasattr(mesh, "_shared_data_cache") hierarchyf, levelf = get_level(Vf.mesh()) @@ -63,10 +60,7 @@ def coarse_node_to_fine_node_map(Vc, Vf): if levelc + increment != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) - key = (entity_dofs_key(Vc.finat_element.entity_dofs()) - + entity_dofs_key(Vf.finat_element.entity_dofs()) - + (levelc, levelf)) - + key = _cache_key(Vc, Vf) cache = mesh._shared_data_cache["hierarchy_coarse_node_to_fine_node_map"] try: return cache[key] @@ -101,7 +95,7 @@ def coarse_cell_to_fine_node_map(Vc, Vf): if levelc + increment != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) - key = (entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf)) + key = _cache_key(Vc, Vf) cache = mesh._shared_data_cache["hierarchy_coarse_cell_to_fine_node_map"] try: return cache[key] @@ -142,11 +136,12 @@ def physical_node_locations(V): # This is a defaultdict, so the first time we access the key we # get a fresh dict for the cache. cache = mesh._geometric_shared_data_cache["hierarchy_physical_node_locations"] - key = element + key = (element, tuple(V.boundary_set)) try: return cache[key] except KeyError: - Vc = firedrake.VectorFunctionSpace(mesh, element) + Vc = V.collapse().reconstruct(element=finat.ufl.VectorElement(element, dim=mesh.geometric_dimension())) + # FIXME: This is unsafe for DG coordinates and CG target spaces. locations = firedrake.assemble(firedrake.Interpolate(firedrake.SpatialCoordinate(mesh), Vc)) return cache.setdefault(key, locations) @@ -171,3 +166,10 @@ def get_level(obj): def has_level(obj): """Does the provided object have level info?""" return hasattr(obj.topological, "__level_info__") + + +def _cache_key(Vc, Vf): + return (entity_dofs_key(Vc.finat_element.entity_dofs()) + + entity_dofs_key(Vf.finat_element.entity_dofs()) + + (Vc.dim(), Vf.dim()) + + (tuple(Vc.boundary_set), tuple(Vf.boundary_set))) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 5fb1fca51a..a39160ac45 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -4,7 +4,6 @@ from firedrake.petsc import PETSc from firedrake.dmhooks import get_function_space, get_appctx from firedrake.ufl_expr import TestFunction, TrialFunction -from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace from ufl import curl, div, HCurl, HDiv, inner, dx from pyop2.utils import as_tuple import numpy @@ -88,13 +87,7 @@ def initialize(self, pc): if B is None: from firedrake.assemble import assemble d = {HCurl: curl, HDiv: div}[sobolev_space] - if V.shape == (): - make_function_space = FunctionSpace - elif len(V.shape) == 1: - make_function_space = VectorFunctionSpace - else: - make_function_space = TensorFunctionSpace - Q = make_function_space(V.mesh(), "DG", degree-1) + Q = V.reconstruct(family="DG", degree=degree-1) b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1)) B = assemble(b, mat_type="matfree") bddcpc.setBDDCDivergenceMat(B.petscmat) @@ -103,7 +96,7 @@ def initialize(self, pc): if gradient is None: from firedrake.preconditioners.fdm import tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad - Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element())) + Q = V.reconstruct(element=curl_to_grad(V.ufl_element())) gradient = tabulate_exterior_derivative(Q, V) corners = get_vertex_dofs(Q) gradient.compose('_elements_corners', corners) @@ -135,7 +128,7 @@ def applyTranspose(self, pc, x, y): def get_vertex_dofs(V): - W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), "vertex")) + W = V.reconstruct(element=restrict(V.ufl_element(), "vertex")) indices = get_restriction_indices(V, W) V.dof_dset.lgmap.apply(indices, result=indices) vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm) diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 4c8e654ce4..5848e0179b 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/firedrake/preconditioners/facet_split.py @@ -4,6 +4,7 @@ from finat.ufl import RestrictedElement, MixedElement, TensorElement, VectorElement from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase +from firedrake.bcs import restricted_function_space import firedrake.dmhooks as dmhooks import numpy @@ -58,6 +59,8 @@ def initialize(self, pc): element = V.ufl_element() elements = [restrict(element, domain) for domain in domains] W = FunctionSpace(V.mesh(), elements[0] if len(elements) == 1 else MixedElement(elements)) + if V.boundary_set: + W = restricted_function_space(W, [V.boundary_set]*len(W)) args = (TestFunction(W), TrialFunction(W)) if len(W) > 1: diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 1b6aa9c9aa..a52fba5aae 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -246,7 +246,7 @@ def allocate_matrix(self, Amat, V, J, bcs, fcp, pmat_type, use_static_condensati elif len(ifacet) == 1: Vfacet = V[ifacet[0]] ebig, = set(unrestrict_element(Vsub.ufl_element()) for Vsub in V) - Vbig = FunctionSpace(V.mesh(), ebig) + Vbig = V.reconstruct(element=ebig) if len(V) > 1: dims = [Vsub.finat_element.space_dimension() for Vsub in V] assert sum(dims) == Vbig.finat_element.space_dimension() @@ -436,9 +436,9 @@ def condense(self, A, J, bcs, fcp, pc_type="icc"): V0 = next((Vi for Vi in V if is_restricted(Vi.finat_element)[0]), None) V1 = next((Vi for Vi in V if is_restricted(Vi.finat_element)[1]), None) if V0 is None: - V0 = FunctionSpace(V.mesh(), restrict_element(self.embedding_element, "interior")) + V0 = V.reconstruct(element=restrict_element(self.embedding_element, "interior")) if V1 is None: - V1 = FunctionSpace(V.mesh(), restrict_element(self.embedding_element, "facet")) + V1 = V.reconstruct(element=restrict_element(self.embedding_element, "facet")) if len(V) == 1: J00 = J(*(t.reconstruct(function_space=V0) for t in J.arguments())) elif len(V) == 2: @@ -446,7 +446,7 @@ def condense(self, A, J, bcs, fcp, pc_type="icc"): ises = V.dof_dset.field_ises Smats[V[0], V[1]] = A.createSubMatrix(ises[0], ises[1]) Smats[V[1], V[0]] = A.createSubMatrix(ises[1], ises[0]) - unindexed = {Vsub: FunctionSpace(Vsub.mesh(), Vsub.ufl_element()) for Vsub in V} + unindexed = {Vsub: Vsub.collapse() for Vsub in V} bcs = tuple(bc.reconstruct(V=unindexed[bc.function_space()], g=0) for bc in bcs) else: raise ValueError("Expecting at most 2 components") @@ -708,7 +708,7 @@ def _element_kernels(self): element_kernel = TripleProductKernel(R1, M, C1) schur_kernel = self.schur_kernel.get(Vrow) if Vrow == Vcol else None if schur_kernel is not None: - V0 = FunctionSpace(Vrow.mesh(), restrict_element(self.embedding_element, "interior")) + V0 = Vrow.collapse().reconstruct(element=restrict_element(self.embedding_element, "interior")) C0 = self.assemble_reference_tensor(V0, sort_interior=True) R0 = self.assemble_reference_tensor(V0, sort_interior=True, transpose=True) element_kernel = schur_kernel(element_kernel, @@ -826,7 +826,7 @@ def set_values(self, A, Vrow, Vcol, mat_type=None): element_kernel = TripleProductKernel(R1, M, C1) schur_kernel = self.schur_kernel.get(Vrow) if on_diag else None if schur_kernel is not None: - V0 = FunctionSpace(Vrow.mesh(), restrict_element(self.embedding_element, "interior")) + V0 = Vrow.collapse().reconstruct(element=restrict_element(self.embedding_element, "interior")) C0 = self.assemble_reference_tensor(V0, sort_interior=True) R0 = self.assemble_reference_tensor(V0, sort_interior=True, transpose=True) element_kernel = schur_kernel(element_kernel, @@ -1462,9 +1462,9 @@ def __init__(self, kernel, name=None): fcp = self.child.fcp args = form.arguments() Q = args[0].function_space() - V = FunctionSpace(Q.mesh(), unrestrict_element(Q.ufl_element())) - V0 = FunctionSpace(Q.mesh(), restrict_element(V.ufl_element(), "interior")) - V1 = FunctionSpace(Q.mesh(), restrict_element(V.ufl_element(), "facet")) + V = Q.reconstruct(element=unrestrict_element(Q.ufl_element())) + V0 = Q.reconstruct(element=restrict_element(V.ufl_element(), "interior")) + V1 = Q.reconstruct(element=restrict_element(V.ufl_element(), "facet")) idofs = PETSc.IS().createBlock(V.block_size, restricted_dofs(V0.finat_element, V.finat_element), comm=comm) fdofs = PETSc.IS().createBlock(V.block_size, restricted_dofs(V1.finat_element, V.finat_element), comm=comm) size = idofs.size + fdofs.size diff --git a/firedrake/preconditioners/hiptmair.py b/firedrake/preconditioners/hiptmair.py index 858a07d0d5..0b9d55b06f 100644 --- a/firedrake/preconditioners/hiptmair.py +++ b/firedrake/preconditioners/hiptmair.py @@ -5,7 +5,6 @@ from firedrake.bcs import DirichletBC from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase -from firedrake.functionspace import FunctionSpace from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.preconditioners.hypre_ams import chop from firedrake.preconditioners.facet_split import restrict @@ -146,7 +145,6 @@ def coarsen(self, pc): a, bcs = self.form(pc) V = a.arguments()[-1].function_space() - mesh = V.mesh() element = V.ufl_element() formdegree = V.finat_element.formdegree if formdegree == 1: @@ -163,7 +161,7 @@ def coarsen(self, pc): if domain: celement = restrict(celement, domain) - coarse_space = FunctionSpace(mesh, celement) + coarse_space = V.reconstruct(element=celement) assert coarse_space.finat_element.formdegree + 1 == formdegree coarse_space_bcs = [bc.reconstruct(V=coarse_space, g=0) for bc in bcs] diff --git a/firedrake/preconditioners/hypre_ads.py b/firedrake/preconditioners/hypre_ads.py index 5affcf8c1f..bf9fbe4b85 100644 --- a/firedrake/preconditioners/hypre_ads.py +++ b/firedrake/preconditioners/hypre_ads.py @@ -1,11 +1,11 @@ import firedrake.assemble as assemble from firedrake.preconditioners.base import PCBase from firedrake.petsc import PETSc -from firedrake.functionspace import FunctionSpace, VectorFunctionSpace from firedrake.ufl_expr import TestFunction from firedrake.dmhooks import get_function_space from firedrake.preconditioners.hypre_ams import chop from firedrake.__future__ import interpolate +from finat.ufl import VectorElement from ufl import grad, curl, SpatialCoordinate from pyop2.utils import as_tuple @@ -26,8 +26,8 @@ def initialize(self, obj): if formdegree != 2 or degree != 1: raise ValueError("Hypre ADS requires lowest order RT elements! (not %s of degree %d)" % (family, degree)) - P1 = FunctionSpace(mesh, "Lagrange", 1) - NC1 = FunctionSpace(mesh, "N1curl" if mesh.ufl_cell().is_simplex() else "NCE", 1) + P1 = V.reconstruct(family="Lagrange", degree=1) + NC1 = V.reconstruct(family="N1curl" if mesh.ufl_cell().is_simplex() else "NCE", degree=1) G_callback = appctx.get("get_gradient", None) if G_callback is None: G = chop(assemble.assemble(interpolate(grad(TestFunction(P1)), NC1)).petscmat) @@ -49,7 +49,7 @@ def initialize(self, obj): pc.setHYPREDiscreteGradient(G) pc.setHYPREDiscreteCurl(C) - VectorP1 = VectorFunctionSpace(mesh, "Lagrange", 1) + VectorP1 = P1.reconstruct(element=VectorElement(P1.ufl_element())) interp = interpolate(SpatialCoordinate(mesh), VectorP1) pc.setCoordinates(assemble.assemble(interp).dat.data_ro.copy()) diff --git a/firedrake/preconditioners/hypre_ams.py b/firedrake/preconditioners/hypre_ams.py index dc610e005d..8445d429aa 100644 --- a/firedrake/preconditioners/hypre_ams.py +++ b/firedrake/preconditioners/hypre_ams.py @@ -1,13 +1,13 @@ import firedrake.assemble as assemble from firedrake.preconditioners.base import PCBase from firedrake.petsc import PETSc -from firedrake.functionspace import FunctionSpace, VectorFunctionSpace from firedrake.ufl_expr import TestFunction from firedrake.dmhooks import get_function_space from firedrake.utils import complex_mode from firedrake_citations import Citations from firedrake import SpatialCoordinate from firedrake.__future__ import interpolate +from finat.ufl import VectorElement from ufl import grad from pyop2.utils import as_tuple @@ -47,7 +47,7 @@ def initialize(self, obj): if formdegree != 1 or degree != 1: raise ValueError("Hypre AMS requires lowest order Nedelec elements! (not %s of degree %d)" % (family, degree)) - P1 = FunctionSpace(mesh, "Lagrange", 1) + P1 = V.reconstruct(family="Lagrange", degree=1) G_callback = appctx.get("get_gradient", None) if G_callback is None: G = chop(assemble.assemble(interpolate(grad(TestFunction(P1)), V)).petscmat) @@ -67,7 +67,7 @@ def initialize(self, obj): if zero_beta: pc.setHYPRESetBetaPoissonMatrix(None) - VectorP1 = VectorFunctionSpace(mesh, "Lagrange", 1) + VectorP1 = P1.reconstruct(element=VectorElement(P1.ufl_element())) interp = interpolate(SpatialCoordinate(mesh), VectorP1) pc.setCoordinates(assemble.assemble(interp).dat.data_ro.copy()) pc.setFromOptions() diff --git a/firedrake/preconditioners/patch.py b/firedrake/preconditioners/patch.py index 6d2bced79f..cdcca9c2f3 100644 --- a/firedrake/preconditioners/patch.py +++ b/firedrake/preconditioners/patch.py @@ -4,7 +4,6 @@ from firedrake.solving_utils import _SNESContext from firedrake.utils import cached_property, complex_mode, IntType from firedrake.dmhooks import get_appctx, push_appctx, pop_appctx -from firedrake.functionspace import FunctionSpace from firedrake.interpolation import Interpolate from collections import namedtuple @@ -654,22 +653,15 @@ def sort_entities(self, dm, axis, dir, ndiv=None, divisions=None): raise RuntimeError("Must either set ndiv or divisions for PlaneSmoother!") mesh = dm.getAttr("__firedrake_mesh__") - ele = mesh.coordinates.function_space().ufl_element() - V = mesh.coordinates.function_space() - if V.finat_element.entity_dofs() == V.finat_element.entity_closure_dofs(): + coordinates = mesh.coordinates + V = coordinates.function_space() + if V.finat_element.is_dg(): # We're using DG or DQ for our coordinates, so we got # a periodic mesh. We need to interpolate to CGk # with access descriptor MAX to define a consistent opinion # about where the vertices are. - CGkele = ele.reconstruct(family="Lagrange") - # Need to supply the actual mesh to the FunctionSpace constructor, - # not its weakref proxy (the variable `mesh`) - # as interpolation fails because they are not hashable - CGk = FunctionSpace(mesh.coordinates.function_space().mesh(), CGkele) - coordinates = Interpolate(mesh.coordinates, CGk, access=op2.MAX) - coordinates = assemble(coordinates) - else: - coordinates = mesh.coordinates + CGk = V.reconstruct(family="Lagrange") + coordinates = assemble(Interpolate(coordinates, CGk, access=op2.MAX)) select = partial(select_entity, dm=dm, exclude="pyop2_ghost") entities = [(p, self.coords(dm, p, coordinates)) for p in diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 48df37d615..04797e95aa 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -181,21 +181,21 @@ def coarsen(self, fdm, comm): assert parent is not None test, trial = fctx.J.arguments() - fV = test.function_space() + fV = trial.function_space() cele = self.coarsen_element(fV.ufl_element()) # Have we already done this? cctx = fctx._coarse if cctx is not None: - cV = cctx.J.arguments()[0].function_space() - if (cV.ufl_element() == cele) and (cV.mesh() == fV.mesh()): + cV = cctx.J.arguments()[1].function_space() + if (cV.ufl_element() == cele) and (cV.mesh() == fV.mesh()) and all(cV_.boundary_set == fV_.boundary_set for cV_, fV_ in zip(cV, fV)): return cV.dm - cV = firedrake.FunctionSpace(fV.mesh(), cele) + cV = fV.reconstruct(element=cele) cdm = cV.dm fproblem = fctx._problem - fu = fproblem.u + fu = fproblem.u_restrict cu = firedrake.Function(cV) fdeg = PMGBase.max_degree(fV.ufl_element()) @@ -370,8 +370,8 @@ def create_transfer(self, mat_type, cctx, fctx, cbcs, fbcs): construct_mat = prolongation_matrix_aij else: raise ValueError("Unknown matrix type") - cV = cctx.J.arguments()[0].function_space() - fV = fctx.J.arguments()[0].function_space() + cV = cctx._problem.u_restrict.function_space() + fV = fctx._problem.u_restrict.function_space() cbcs = tuple(cctx._problem.bcs) if cbcs else tuple() fbcs = tuple(fctx._problem.bcs) if fbcs else tuple() return cache.setdefault(key, construct_mat(cV, fV, cbcs, fbcs)) @@ -1179,7 +1179,7 @@ def make_permutation_code(V, vshape, pshape, t_in, t_out, array_name): def reference_value_space(V): element = finat.ufl.WithMapping(V.ufl_element(), mapping="identity") - return firedrake.FunctionSpace(V.mesh(), element) + return V.collapse().reconstruct(element=element) class StandaloneInterpolationMatrix(object): @@ -1206,13 +1206,13 @@ def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs): self.Vf = reference_value_space(self.Vf) self.uc = firedrake.Function(self.Vc, val=self.uc.dat) self.uf = firedrake.Function(self.Vf, val=self.uf.dat) - self.Vc_bcs = [bc.reconstruct(V=self.Vc) for bc in self.Vc_bcs] - self.Vf_bcs = [bc.reconstruct(V=self.Vf) for bc in self.Vf_bcs] + self.Vc_bcs = [bc.reconstruct(V=self.Vc, g=0) for bc in self.Vc_bcs] + self.Vf_bcs = [bc.reconstruct(V=self.Vf, g=0) for bc in self.Vf_bcs] def work_function(self, V): if isinstance(V, firedrake.Function): return V - key = (V.ufl_element(), V.mesh()) + key = (V.ufl_element(), V.mesh(), tuple(V.boundary_set)) try: return self._cache_work[key] except KeyError: @@ -1337,17 +1337,14 @@ def make_blas_kernels(self, Vf, Vc): restrict = [""]*5 # get embedding element for Vf with identity mapping and collocated vector component DOFs try: - qelem = felem - if qelem.mapping() != "identity": - qelem = qelem.reconstruct(mapping="identity") - Qf = Vf if qelem == felem else firedrake.FunctionSpace(Vf.mesh(), qelem) + Qf = Vf if felem.mapping() == "identity" else Vf.reconstruct(mapping="identity") mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1") in_place_mapping = True except Exception: qelem = finat.ufl.FiniteElement("DQ", cell=felem.cell, degree=PMGBase.max_degree(felem)) if Vf.value_shape: qelem = finat.ufl.TensorElement(qelem, shape=Vf.value_shape, symmetry=felem.symmetry()) - Qf = firedrake.FunctionSpace(Vf.mesh(), qelem) + Qf = Vf.reconstruct(element=qelem) mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1") qshape = (Qf.block_size, Qf.finat_element.space_dimension()) diff --git a/pyop2/types/dat.py b/pyop2/types/dat.py index 45900effef..7dcd08e589 100644 --- a/pyop2/types/dat.py +++ b/pyop2/types/dat.py @@ -105,7 +105,6 @@ def _wrapper_cache_key_(self): @utils.validate_in(('access', _modes, ex.ModeValueError)) def __call__(self, access, path=None): from pyop2.parloop import DatLegacyArg - if conf.configuration["type_check"] and path and path.toset != self.dataset.set: raise ex.MapValueError("To Set of Map does not match Set of Dat.") return DatLegacyArg(self, path, access) diff --git a/tests/firedrake/multigrid/test_p_multigrid.py b/tests/firedrake/multigrid/test_p_multigrid.py index 39437b0204..ea26f1c44a 100644 --- a/tests/firedrake/multigrid/test_p_multigrid.py +++ b/tests/firedrake/multigrid/test_p_multigrid.py @@ -137,7 +137,8 @@ def mat_type(request): return request.param -def test_p_multigrid_scalar(mesh, mat_type): +@pytest.mark.parametrize("restrict", [False, True], ids=("unrestrict", "restrict")) +def test_p_multigrid_scalar(mesh, mat_type, restrict): V = FunctionSpace(mesh, "CG", 4) u = Function(V) @@ -175,7 +176,7 @@ def test_p_multigrid_scalar(mesh, mat_type): "pmg_mg_coarse_mg_coarse_ksp_monitor": None, "pmg_mg_coarse_mg_coarse_pc_type": "gamg", "pmg_mg_coarse_mg_coarse_pc_gamg_threshold": 0} - problem = NonlinearVariationalProblem(F, u, bcs) + problem = NonlinearVariationalProblem(F, u, bcs, restrict=restrict) solver = NonlinearVariationalSolver(problem, solver_parameters=sp) solver.solve() diff --git a/tests/firedrake/regression/test_restricted_function_space.py b/tests/firedrake/regression/test_restricted_function_space.py index 8d8a19a56d..ddbc876ddc 100644 --- a/tests/firedrake/regression/test_restricted_function_space.py +++ b/tests/firedrake/regression/test_restricted_function_space.py @@ -362,3 +362,97 @@ def test_restricted_function_space_extrusion_stokes(ncells): # -- Actually, the ordering is the same. assert np.allclose(sol_res.subfunctions[0].dat.data_ro_with_halos, sol.subfunctions[0].dat.data_ro_with_halos) assert np.allclose(sol_res.subfunctions[1].dat.data_ro_with_halos, sol.subfunctions[1].dat.data_ro_with_halos) + + +@pytest.mark.parametrize("names", [(None, None), (None, "name1"), ("name0", "name1")]) +def test_restrict_fieldsplit(names): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 1, name=names[0]) + Q = FunctionSpace(mesh, "CG", 2, name=names[1]) + Z = V * Q + + z = Function(Z) + test = TestFunction(Z) + z_exact = Constant([1, -1]) + + F = inner(z - z_exact, test) * dx + bcs = [DirichletBC(Z.sub(i), z_exact[i], (i+1, i+3)) for i in range(len(Z))] + + problem = NonlinearVariationalProblem(F, z, bcs=bcs, restrict=True) + solver = NonlinearVariationalSolver(problem, solver_parameters={ + "snes_type": "ksponly", + "ksp_type": "preonly", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + f"fieldsplit_{names[0] or 0}_pc_type": "lu", + f"fieldsplit_{names[1] or 1}_pc_type": "lu"}, + options_prefix="") + solver.solve() + + # Test prefixes for the restricted spaces + pc = solver.snes.ksp.pc + for field, ksp in enumerate(pc.getFieldSplitSubKSP()): + name = Z[field].name or field + assert ksp.getOptionsPrefix() == f"fieldsplit_{name}_" + + assert errornorm(z_exact[0], z.subfunctions[0]) < 1E-10 + assert errornorm(z_exact[1], z.subfunctions[1]) < 1E-10 + + +def test_restrict_python_pc(): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 1) + u = Function(V) + test = TestFunction(V) + + x, y = SpatialCoordinate(mesh) + u_exact = x + y + g = Function(V).interpolate(u_exact) + + F = inner(u - u_exact, test) * dx + bcs = [DirichletBC(V, g, 1), DirichletBC(V, u_exact, 2)] + + problem = NonlinearVariationalProblem(F, u, bcs=bcs, restrict=True) + solver = NonlinearVariationalSolver(problem, solver_parameters={ + "snes_type": "ksponly", + "mat_type": "matfree", + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled_pc_type": "lu"}) + solver.solve() + + assert errornorm(u_exact, u) < 1E-10 + + +def test_restrict_multigrid(): + base = UnitSquareMesh(2, 2) + refine = 2 + mh = MeshHierarchy(base, refine) + mesh = mh[-1] + + V = FunctionSpace(mesh, "CG", 1) + u = Function(V) + test = TestFunction(V) + + x, y = SpatialCoordinate(mesh) + u_exact = x + y + g = Function(V).interpolate(u_exact) + + F = inner(grad(u - u_exact), grad(test)) * dx + bcs = [DirichletBC(V, g, 1), DirichletBC(V, u_exact, 2)] + + problem = NonlinearVariationalProblem(F, u, bcs=bcs, restrict=True) + solver = NonlinearVariationalSolver(problem, solver_parameters={ + "snes_type": "ksponly", + "ksp_type": "cg", + "ksp_rtol": 1E-10, + "ksp_max_it": 10, + "ksp_monitor": None, + "pc_type": "mg", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_pc_type": "jacobi", + "mg_coarse_pc_type": "lu"}) + solver.solve() + + assert errornorm(u_exact, u) < 1E-10