diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 815ae360a2..1d2aa67aae 100644 --- a/firedrake/formmanipulation.py +++ b/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: diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 0fdfbe8e20..65935e5faf 100644 --- a/firedrake/functionspaceimpl.py +++ b/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): diff --git a/firedrake/slate/slate.py b/firedrake/slate/slate.py index 58e281368f..bb6909128f 100644 --- a/firedrake/slate/slate.py +++ b/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())) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 789a6f1880..c6444f0daa 100644 --- a/firedrake/solving_utils.py +++ b/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: