WIP: Allow indexing in replace_test_function and replace_trial_function maps.

[JHopeCollins]

This will allow passing an idx argument to replace_test_function and replace_trial_function in the same way as replace_subject, which will help when dealing with mixed spaces.

[JHopeCollins]

To allow adding the idx argument to replace_test_function and replace_trial_function I factored out the logic in replace_subject for building the dictionary for ufl.replace so that it could be used for all three replace_* functions.
This logic has two components:

1. Which type combinations are allowable for the new values and the old values being replaced?
2. How to build the dictionary for ufl.replace from a given combination of old/new values?

When I was refactoring I ran into into issues with having to check for extra (valid) type combinations between the new/old values. But 1. is really just asking "will ufl.replace accept this dictionary?", so the type-checking can be left to UFL, and we only need to deal with 2. i.e. building the dictionary. This is also preferable in the case that UFL/Firedrake introduce new types in the future that we would otherwise have to deal with.
The only type-checking for 2. is to see whether the old/new values are indexable (i.e. tuple or member of a MixedFunctionSpace), which is much more manageable. This does mean that the dictionary may contain invalid key-value combinations, but the dictionary is never passed back to user-land before being passed to ufl.replace so I don't think this is an issue.

[JHopeCollins]

I merged main into this branch and I'm now seeing errors when trying to use replace_subject with split components of a mixed space. The errors come from the ufl.replace call added in #346 to deal with the perp operator.

Failing example:

import firedrake as fd
import gusto

### set up gusto equation

mesh = fd.IcosahedralSphereMesh(radius=1, refinement_level=2, degree=1)
x, y, z = fd.SpatialCoordinate(mesh)

swe_parameters = gusto.ShallowWaterParameters(H=1, g=1, Omega=1)

dt = 1
domain = gusto.Domain(mesh, dt, 'BDM', degree=1)

eqn = gusto.ShallowWaterEquations(domain, swe_parameters, fexpr=z)

### get mass and stiffness forms

from gusto.labels import (replace_subject, replace_test_function, time_derivative)
from gusto.fml.form_manipulation_labelling import all_terms, drop

M = eqn.residual.label_map(lambda t: t.has_label(time_derivative),
                           map_if_false=drop)

K = eqn.residual.label_map(lambda t: t.has_label(time_derivative),
                           map_if_true=drop)
                                                                                                                                                                                                  
### replacing arguments

W = eqn.function_space

w = fd.Function(W)
u, h = fd.split(w)
v, q = fd.TestFunctions(W)

# test functions

R = eqn.residual
F = R.label_map(all_terms, replace_test_function(v, idx=0)) # works
F = R.label_map(all_terms, replace_test_function(q, idx=1)) # works

# subject

F = R.label_map(all_terms, replace_subject(w)) # works

F = R.label_map(all_terms, replace_subject(u, idx=0)) # breaks 1
F = R.label_map(all_terms, replace_subject(h, idx=1)) # breaks 2

Replacing the indexed test functions works (yay!), but replacing the split subject does not. Replacing the unsplit subject works as before. The last two lines produce the following errors, but don't error if the perp if-block in replace_subject is commented out.

Error for break 1:

UFL:ERROR Function splitting failed to extract components for whole intended range. Something is wrong.
Traceback (most recent call last):
  File "gusto_extract_equation_form.py", line 41, in <module>
    F = R.label_map(all_terms, replace_subject(u, idx=0)) # breaks 1
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/fml/form_manipulation_labelling.py", line 286, in label_map
    functools.reduce(operator.add,
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/fml/form_manipulation_labelling.py", line 288, in <genexpr>
    (map_if_true(t) if term_filter(t) else
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/labels.py", line 191, in repl
    new_form = ufl.replace(new_form, {split(new_subj)[0]: perp_function(split(new_subj)[0])})
  File "/home/jhc/icl_ra/firedrake/src/ufl/ufl/split_functions.py", line 109, in split
    error("Function splitting failed to extract components for whole intended range. Something is wrong.")
  File "/home/jhc/icl_ra/firedrake/src/ufl/ufl/log.py", line 158, in error
    raise self._exception_type(self._format_raw(*message))
ufl.log.UFLException: Function splitting failed to extract components for whole intended range. Something is wrong.

Error for break 2:

UFL:ERROR Cross product requires arguments of rank 1, got <CellNormal id=23437979577088> and <Indexed id=23438038609984>.
Traceback (most recent call last):
  File "gusto_extract_equation_form.py", line 41, in <module>
    F = R.label_map(all_terms, replace_subject(h, idx=1)) # breaks 2
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/fml/form_manipulation_labelling.py", line 286, in label_map
    functools.reduce(operator.add,
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/fml/form_manipulation_labelling.py", line 288, in <genexpr>
    (map_if_true(t) if term_filter(t) else
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/labels.py", line 191, in repl
    new_form = ufl.replace(new_form, {split(new_subj)[0]: perp_function(split(new_subj)[0])})
  File "/home/jhc/icl_ra/firedrake/src/gusto/gusto/domain.py", line 81, in <lambda>
    self.perp = lambda u: cross(outward_normals, u)
  File "/home/jhc/icl_ra/firedrake/src/ufl/ufl/operators.py", line 215, in cross
    return Cross(a, b)
  File "/home/jhc/icl_ra/firedrake/src/ufl/ufl/tensoralgebra.py", line 230, in __new__
    error("Cross product requires arguments of rank 1, got %s and %s." % (
  File "/home/jhc/icl_ra/firedrake/src/ufl/ufl/log.py", line 158, in error
    raise self._exception_type(self._format_raw(*message))
ufl.log.UFLException: Cross product requires arguments of rank 1, got <CellNormal id=23437979577088> and <Indexed id=23438038609984>.

I'm not entirely sure what the fix in #346 is doing or why it is necessary, so I don't really know how to go about fixing this. @jshipton @tommbendall any ideas?