MixedFunctionSpace: interpolate and restrict

firedrake/eigensolver.py

@@ -1,6 +1,5 @@
 """Specify and solve finite element eigenproblems."""
 from firedrake.assemble import assemble
-from firedrake.bcs import DirichletBC
 from firedrake.function import Function
 from firedrake.functionspace import RestrictedFunctionSpace
 from firedrake.ufl_expr import TrialFunction, TestFunction
@@ -71,12 +70,12 @@ def __init__(self, A, M=None, bcs=None, bc_shift=0.0, restrict=True):
             M = inner(u, v) * dx
 
         if restrict and bcs:  # assumed u and v are in the same space here
-            V_res = RestrictedFunctionSpace(self.output_space, boundary_set=set([bc.sub_domain for bc in bcs]))
+            V_res = RestrictedFunctionSpace(self.output_space, boundary_set=set(bc.sub_domain for bc in bcs))
             u_res = TrialFunction(V_res)
             v_res = TestFunction(V_res)
             self.M = replace(M, {u: u_res, v: v_res})
             self.A = replace(A, {u: u_res, v: v_res})
-            self.bcs = [DirichletBC(V_res, bc.function_arg, bc.sub_domain) for bc in bcs]
+            self.bcs = [bc.reconstruct(V=V_res) for bc in bcs]
             self.restricted_space = V_res
         else:
             self.A = A  # LHS

firedrake/functionspaceimpl.py

@@ -860,17 +860,35 @@ class RestrictedFunctionSpace(FunctionSpace):
     output of the solver.
 
     :arg function_space: The :class:`FunctionSpace` to restrict.
+    :kwarg boundary_set: An iterable of DirichletBCs or a set of subdomains on
+        which a DirichletBC will be applied.
     :kwarg name: An optional name for this :class:`RestrictedFunctionSpace`,
         useful for later identification.
-    :kwarg boundary_set: A set of subdomains on which a DirichletBC will be
-        applied.
 
     Notes
     -----
     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 __init__(self, function_space, name=None, boundary_set=frozenset()):
+    def __new__(cls, *args, boundary_set=frozenset(), name=None, **kwargs):
+        if cls is not RestrictedFunctionSpace:
+            return super().__new__(cls)
+        function_space = args[0]
+        if len(function_space) > 1:
+            return MixedFunctionSpace([cls(Vsub, boundary_set=boundary_set)
+                                       for Vsub in function_space], name=name)
+        return super().__new__(cls)
+
+    def __init__(self, function_space, boundary_set=frozenset(), name=None):
+        if all(hasattr(bc, "sub_domain") for bc in boundary_set):
+            bcs = boundary_set
+            boundary_set = []
+            for bc in bcs:
+                if bc.function_space() == function_space:
+                    if type(bc.sub_domain) in {str, int}:
+                        boundary_set.append(bc.sub_domain)
+                    else:
+                        boundary_set.extend(bc.sub_domain)
         label = ""
         for boundary_domain in boundary_set:
             label += str(boundary_domain)
@@ -884,8 +902,7 @@ def __init__(self, function_space, name=None, boundary_set=frozenset()):
                                                      label=self._label)
         self.function_space = function_space
         self.name = name or (function_space.name or "Restricted" + "_"
-                             + "_".join(sorted(
-                                        [str(i) for i in self.boundary_set])))
+                             + "_".join(sorted(map(str, self.boundary_set))))
 
     def set_shared_data(self):
         sdata = get_shared_data(self._mesh, self.ufl_element(), self.boundary_set)

tests/regression/test_restricted_function_space.py

@@ -168,17 +168,27 @@ def test_restricted_function_space_coord_change(j):
     new_mesh = Mesh(Function(V).interpolate(as_vector([x, y])))
     new_V = FunctionSpace(new_mesh, "CG", j)
     bc = DirichletBC(new_V, 0, 1)
-    new_V_restricted = RestrictedFunctionSpace(new_V, name="Restricted", boundary_set=[1])
+    new_V_restricted = RestrictedFunctionSpace(new_V, boundary_set=[1], name="Restricted")
 
     compare_function_space_assembly(new_V, new_V_restricted, [bc])
 
 
+def test_restricted_mixed_spaces():
+    mesh = UnitSquareMesh(1, 1)
+    V = FunctionSpace(mesh, "RT", 1)
+    Q = FunctionSpace(mesh, "DG", 0)
+    Z = V * Q
+    bcs = [DirichletBC(Z.sub(0), 0, [1])]
+    Z_restricted = RestrictedFunctionSpace(Z, bcs)
+    compare_function_space_assembly(Z, Z_restricted, bcs)
+
+
 @pytest.mark.parametrize(["i", "j"], [(1, 0), (2, 0), (2, 1)])
-def test_restricted_mixed_spaces(i, j):
+def test_poisson_mixed_restricted_spaces(i, j):
     mesh = UnitSquareMesh(1, 1)
     DG = FunctionSpace(mesh, "DG", j)
     CG = VectorFunctionSpace(mesh, "CG", i)
-    CG_res = RestrictedFunctionSpace(CG, "Restricted", boundary_set=[4])
+    CG_res = RestrictedFunctionSpace(CG, boundary_set=[4], name="Restricted")
     W = CG * DG
     W_res = CG_res * DG
     bc = DirichletBC(W.sub(0), 0, 4)