PC: FunctionSpace() -> V.reconstruct()

Author: pbrubeck

firedrake/functionspaceimpl.py

@@ -889,7 +889,6 @@ class RestrictedFunctionSpace(FunctionSpace):
     def __new__(cls, function_space, boundary_set=frozenset(), name=None):
         mesh = function_space.mesh()
         if mesh is not mesh.topology:
-            # Restrict WithGeometry and return a WithGeometry
             V = RestrictedFunctionSpace(function_space.topological,
                                         boundary_set=boundary_set, name=name)
             return type(function_space).create(V, mesh)

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)

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:

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,17 +436,17 @@ 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:
             J00 = ExtractSubBlock().split(J, argument_indices=(V0.index, V0.index))
             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

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]
 

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())
 

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()

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,18 @@ 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