TSFC: support same-mesh ufl.Interpolate in Form kernels

.github/actions/install/action.yml

@@ -150,11 +150,14 @@ runs:
           EXTRA_PIP_FLAGS='--no-build-isolation'
         elif [ ${{ inputs.base_ref }} = 'release' ]; then
           EXTRA_PIP_FLAGS=''
+          export PIP_BUILD_CONSTRAINT=constraints.txt
         else
           echo "Unrecognised 'base_ref' input: '${{ inputs.base_ref }}"
           exit 1
         fi
 
+        : # DROP BEFORE MERGE
+        pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/dual-enriched-again
         pip install --verbose $EXTRA_PIP_FLAGS \
           --no-binary h5py \
           --extra-index-url https://download.pytorch.org/whl/cpu \

.github/workflows/core.yml

@@ -564,6 +564,9 @@ jobs:
       - name: Install Firedrake
         id: install
         run: |
+          : # DROP BEFORE MERGE
+          pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/dual-enriched-again
+          pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/ufl.git@pbrubeck/interpolate-reference-lowering
           pip install --verbose -r ./firedrake-repo/requirements-build.txt
           CC=mpicc CXX=mpicxx \
             pip install --verbose --no-build-isolation './firedrake-repo[docs]'

docs/source/conf.py

@@ -167,6 +167,7 @@
     ('py:class', 'pyop2.caching.Cached'),
     ('py:class', 'pyop2.op2.Kernel'),
     ('py:class', 'pyop2.types.mat.Mat'),
+    ('py:class', 'pyop2.types.access.Access'),
     # Ignore mission docs from Firedrake internal "private" code
     # Any "Base" class eg:
     #   firedrake.adjoint.checkpointing.CheckpointBase

docs/source/interpolation.rst

@@ -458,7 +458,8 @@ each block given by
    \end{pmatrix}
 
 The off-diagonal blocks are zero since the dofs are applied component-wise. Firedrake's form
-compiler recognises this and avoids assembling the zero blocks.
+compiler recognises this and avoids assembling the zero blocks, which the nest
+still allocates.
 
 We can assemble more general interpolation matrices between mixed function spaces by interpolating
 vector expressions with arguments. For example, by doing

firedrake/__init__.py

@@ -70,6 +70,7 @@ def init_petsc():
     FiredrakeException, ConvergenceError, MismatchingDomainError,
     VertexOnlyMeshMissingPointsError, DofNotDefinedError, DofTypeError,
     SerialExecutionOnlyError, PointNotInDomainError,
+    MismatchingFunctionSpaceError,
 )
 from firedrake.function import (  # noqa: F401
     Function, CoordinatelessFunction, PointEvaluator

firedrake/bcs.py

@@ -98,6 +98,14 @@
             V = V.parent
         return V
 
+    @property
+    def parent_function_space(self):
+        """The top-level function space, walking up through indexed or component subspaces."""
+        V = self._function_space
+        while V.parent is not None:
+            V = V.parent
+        return V
+
     @cached_property
     def domain_args(self):
         r"""The sub_domain the BC applies to."""

firedrake/exceptions.py

@@ -43,6 +43,12 @@ def __str__(self):
         )
 
 
+class MismatchingFunctionSpaceError(FiredrakeException):
+    """Raised when a function space does not match the one expected, such as a
+    boundary condition applied to a form whose arguments live elsewhere.
+    """
+
+
 class NonUniqueMeshSequenceError(FiredrakeException):
     """Raised when calling `.unique()` on a MeshSequence which contains
     non-unique meshes.

firedrake/formmanipulation.py

@@ -1,9 +1,10 @@
 
 import numpy
 import collections
+from collections.abc import Iterable
 
 from ufl import as_tensor, as_vector, split
-from ufl.classes import Form, Zero, FixedIndex, ListTensor, ZeroBaseForm
+from ufl.classes import Expr, Form, Interpolate, Zero, FixedIndex, ListTensor, ZeroBaseForm
 from ufl.algorithms.map_integrands import map_integrand_dags
 from ufl.algorithms import expand_derivatives
 from ufl.corealg.map_dag import MultiFunction, map_expr_dags
@@ -13,6 +14,7 @@
 
 from firedrake.petsc import PETSc
 from firedrake.functionspace import MixedFunctionSpace
+from firedrake.functionspaceimpl import WithGeometry
 from firedrake.cofunction import Cofunction
 from firedrake.ufl_expr import Coargument
 
@@ -30,6 +32,12 @@ class ExtractSubBlock(MultiFunction):
 
     """Extract a sub-block from a form."""
 
+    def __init__(self):
+        super().__init__()
+        self._arg_cache = {}
+        self.blocks = {}
+        self._splitting_interpolate = False
+
     class IndexInliner(MultiFunction):
         """Inline fixed index of list tensors"""
         expr = MultiFunction.reuse_if_untouched
@@ -82,6 +90,9 @@ def split(self, form, argument_indices):
         args = form.arguments()
         self._arg_cache = {}
         self.blocks = dict(enumerate(map(as_tuple, argument_indices)))
+        # An outermost Interpolate splits into a smaller Interpolate, while one
+        # inside an integrand must keep the value shape its neighbours expect.
+        self._splitting_interpolate = isinstance(form, Interpolate)
         if len(args) == 0:
             # Functional can't be split
             return form
@@ -229,9 +240,38 @@ def matrix(self, o):
     def zero_base_form(self, o):
         return ZeroBaseForm(tuple(map(self, o.arguments())))
 
+    def _zero_interpolate(self, o: Interpolate) -> ZeroBaseForm | Zero:
+        """Result of an Interpolate whose operand or target block is Zero."""
+        if self._splitting_interpolate:
+            return self(ZeroBaseForm(o.arguments()))
+        return Zero(o.ufl_shape)
+
+    @staticmethod
+    def _select_components(V: WithGeometry, indices: tuple, operand: Expr) -> list:
+        """Flatten the sub-blocks of ``operand`` whose subspace is in ``indices``."""
+        components = []
+        cur = 0
+        for i, Vi in enumerate(V):
+            if i in indices:
+                components.extend(operand[k] for k in range(cur, cur + Vi.value_size))
+            cur += Vi.value_size
+        return components
+
+    @staticmethod
+    def _embed_components(V: WithGeometry, indices: tuple, values: Iterable) -> list:
+        """Embed ``values`` into V's full shape, zero outside ``indices``."""
+        values = iter(values)
+        components = []
+        for i, Vi in enumerate(V):
+            if i in indices:
+                components.extend(next(values) for _ in range(Vi.value_size))
+            else:
+                components.extend(Zero() for _ in range(Vi.value_size))
+        return components
+
     def interpolate(self, o, operand):
         if isinstance(operand, Zero):
-            return self(ZeroBaseForm(o.arguments()))
+            return self._zero_interpolate(o)
 
         dual_arg, _ = o.argument_slots()
         if len(dual_arg.arguments()) == 1 or len(dual_arg.arguments()[-1].function_space()) == 1:
@@ -249,18 +289,21 @@ def interpolate(self, o, operand):
         W = sub_dual_arg.function_space()
 
         # Unflatten the expression into the target shape
-        cur = 0
-        components = []
-        for i, Vi in enumerate(V):
-            if i in indices:
-                components.extend(operand[i] for i in range(cur, cur+Vi.value_size))
-            cur += Vi.value_size
-
+        components = self._select_components(V, indices, operand)
         operand = as_tensor(numpy.reshape(components, W.value_shape))
         if isinstance(operand, Zero):
-            return self(ZeroBaseForm(o.arguments()))
-
-        return o._ufl_expr_reconstruct_(operand, sub_dual_arg)
+            return self._zero_interpolate(o)
+
+        interpolation = o._ufl_expr_reconstruct_(operand, sub_dual_arg)
+        if self._splitting_interpolate:
+            return interpolation
+
+        # Inside an integrand the block is one part of a wider expression, so
+        # pad it back out to V's value shape with zeros.
+        interpolation_components = ((interpolation[j] for j in numpy.ndindex(interpolation.ufl_shape))
+                                    if interpolation.ufl_shape else iter((interpolation,)))
+        components = self._embed_components(V, indices, interpolation_components)
+        return as_tensor(numpy.reshape(components, V.value_shape))
 
 
 SplitForm = collections.namedtuple("SplitForm", ["indices", "form"])