Feature request: interpolate fails on indexed component functions

[pbrubeck]

Describe the bug
We are able to assign individual components of a vector function u extracted as u.sub(i), but interpolate fails on individual components.

Steps to Reproduce
Steps to reproduce the behavior:

from firedrake import *
mesh = UnitSquareMesh(1, 1)
x = mesh.coordinates
V = x.function_space()
u = Function(V)

u.sub(0).assign(x.sub(0))  # works
u.sub(0).interpolate(x.sub(0))  # fails

Expected behavior
Interpolation of indexed components should be supported.

Error message

File petsc4py/PETSc/Log.pyx:188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File petsc4py/PETSc/Log.pyx:189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/function.py:413, in Function.interpolate(self, expression, subset, allow_missing_dofs, default_missing_val, ad_block_tag)
    408 V = self.function_space()
    409 interp = interpolation.Interpolate(expression, V,
    410                                    subset=subset,
    411                                    allow_missing_dofs=allow_missing_dofs,
    412                                    default_missing_val=default_missing_val)
--> 413 return assemble(interp, tensor=self, ad_block_tag=ad_block_tag)

File petsc4py/PETSc/Log.pyx:188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File petsc4py/PETSc/Log.pyx:189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/adjoint_utils/assembly.py:30, in annotate_assemble.<locals>.wrapper(form, *args, **kwargs)
     28         form = BaseFormAssembler.preprocess_base_form(form)
     29         kwargs['is_base_form_preprocessed'] = True
---> 30     output = assemble(form, *args, **kwargs)
     32 from firedrake.function import Function
     33 from firedrake.cofunction import Cofunction

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:133, in assemble(expr, *args, **kwargs)
    131     raise RuntimeError(f"Got unexpected args: {args}")
    132 tensor = kwargs.pop("tensor", None)
--> 133 return get_assembler(expr, *args, **kwargs).assemble(tensor=tensor)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:381, in BaseFormAssembler.assemble(self, tensor)
    379 # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly.
    380 visited = {}
--> 381 result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited)
    383 if tensor:
    384     BaseFormAssembler.update_tensor(result, tensor)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:603, in BaseFormAssembler.base_form_postorder_traversal(expr, visitor, visited)
    601         stack.extend(unvisited_children)
    602     else:
--> 603         visited[e] = visitor(e, *(visited[arg] for arg in operands))
    605 return visited[expr]

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:377, in BaseFormAssembler.assemble.<locals>.visitor(e, *operands)
    375 def visitor(e, *operands):
    376     t = tensor if e is self._form else None
--> 377     return self.base_form_assembly_visitor(e, t, *operands)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/assemble.py:550, in BaseFormAssembler.base_form_assembly_visitor(self, expr, tensor, *args)
    548     if tensor is None:
    549         return interpolator._interpolate(default_missing_val=default_missing_val)
--> 550     return firedrake.Interpolator(expression, tensor, **interp_data)._interpolate(default_missing_val=default_missing_val)
    551 elif rank == 2:
    552     res = tensor.petscmat if tensor else PETSc.Mat()

File /scratch/brubeckmarti/firedrake/src/pyadjoint/pyadjoint/tape.py:110, in no_annotations.<locals>.wrapper(*args, **kwargs)
    107 @wraps(function)
    108 def wrapper(*args, **kwargs):
    109     with stop_annotating():
--> 110         return function(*args, **kwargs)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/interpolation.py:818, in SameMeshInterpolator.__init__(self, expr, V, subset, freeze_expr, access, bcs, **kwargs)
    816 super().__init__(expr, V, subset, freeze_expr, access, bcs)
    817 try:
--> 818     self.callable, arguments = make_interpolator(expr, V, subset, access, bcs=bcs)
    819 except FIAT.hdiv_trace.TraceError:
    820     raise NotImplementedError("Can't interpolate onto traces sorry")

File petsc4py/PETSc/Log.pyx:188, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File petsc4py/PETSc/Log.pyx:189, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/interpolation.py:999, in make_interpolator(expr, V, subset, access, bcs)
    996 if len(V) > 1:
    997     raise NotImplementedError(
    998         "UFL expressions for mixed functions are not yet supported.")
--> 999 loops.extend(_interpolator(V, tensor, expr, subset, arguments, access, bcs=bcs))
   1000 if bcs and len(arguments) == 0:
   1001     loops.extend([partial(bc.apply, f) for bc in bcs])

File <decorator-gen-49>:2, in _interpolator(V, tensor, expr, subset, arguments, access, bcs)

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/utils.py:89, in known_pyop2_safe.<locals>.wrapper(f, *args, **kwargs)
     87 opts["type_check"] = safe
     88 try:
---> 89     return f(*args, **kwargs)
     90 finally:
     91     opts["type_check"] = check

File /scratch/brubeckmarti/firedrake/src/firedrake/firedrake/interpolation.py:1130, in _interpolator(V, tensor, expr, subset, arguments, access, bcs)
   1128 assert access == op2.WRITE  # Other access descriptors not done for Matrices.
   1129 rows_map = V.cell_node_map()
-> 1130 Vcol = arguments[0].function_space()
   1131 if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
   1132     columns_map = Vcol.cell_node_map()

IndexError: list index out of range