Revert WithGeometry.__getitem__

firedrake/formmanipulation.py

@@ -12,9 +12,19 @@
 from pyop2.utils import as_tuple
 
 from firedrake.petsc import PETSc
+from firedrake.functionspace import MixedFunctionSpace
 from firedrake.cofunction import Cofunction
 
 
+def subspace(V, indices):
+    """Construct a collapsed subspace using components from V."""
+    if len(indices) == 1:
+        W = V[indices[0]]
+    else:
+        W = MixedFunctionSpace([V[i] for i in indices])
+    return W.collapse()
+
+
 class ExtractSubBlock(MultiFunction):
 
     """Extract a sub-block from a form."""
@@ -40,7 +50,7 @@ def indexed(self, o, child, multiindex):
     index_inliner = IndexInliner()
 
     def _subspace_argument(self, a):
-        return type(a)(a.function_space()[self.blocks[a.number()]].collapse(),
+        return type(a)(subspace(a.function_space(), self.blocks[a.number()]),
                        a.number(), part=a.part())
 
     @PETSc.Log.EventDecorator()
@@ -131,7 +141,7 @@ def cofunction(self, o):
 
         # We only need the test space for Cofunction
         indices = self.blocks[0]
-        W = V[indices].collapse()
+        W = subspace(V, indices)
         if len(W) == 1:
             return Cofunction(W, val=o.dat[indices[0]])
         else:

firedrake/functionspaceimpl.py

@@ -321,14 +321,6 @@ def __iter__(self):
         return iter(self.subfunctions)
 
     def __getitem__(self, i):
-        from firedrake.functionspace import MixedFunctionSpace
-        if isinstance(i, (tuple, list)):
-            # Return a subspace
-            if len(i) == 1:
-                return self[i[0]]
-            else:
-                return MixedFunctionSpace([self[isub] for isub in i])
-
         return self.subfunctions[i]
 
     def __mul__(self, other):

firedrake/slate/slate.py

@@ -21,7 +21,7 @@
 from ufl import Constant
 from ufl.coefficient import BaseCoefficient
 
-from firedrake.formmanipulation import ExtractSubBlock
+from firedrake.formmanipulation import ExtractSubBlock, subspace
 from firedrake.function import Function, Cofunction
 from firedrake.ufl_expr import TestFunction
 from firedrake.utils import cached_property, unique
@@ -674,7 +674,7 @@ def _split_arguments(self):
         spaces determined by the indices.
         """
         tensor, = self.operands
-        return tuple(type(a)(a.function_space()[self._blocks[i]].collapse(),
+        return tuple(type(a)(subspace(a.function_space(), self._blocks[i]),
                              a.number(), part=a.part())
                      for i, a in enumerate(tensor.arguments()))
 

firedrake/solving_utils.py

@@ -332,15 +332,13 @@ def split(self, fields):
                 subu = function.Function(V, val=val)
                 # Split it apart to shove in the form.
                 subsplit = split(subu)
-            # Permutation from field indexing to indexing of pieces
-            field_renumbering = {f: i for i, f in enumerate(field)}
             vec = []
             for i, u in enumerate(us):
                 if i in field:
                     # If this is a field we're keeping, get it from
                     # the new function. Otherwise just point to the
                     # old data.
-                    u = subsplit[field_renumbering[i]]
+                    u = subsplit[field.index(i)]
                 if u.ufl_shape == ():
                     vec.append(u)
                 else: