MixedFunctionSpace: interpolate and restrict

[pbrubeck]

Description

This PR enables solving mixed probelms with restrict=True by introducing two main features:

1. We implemented interpolation of flat UFL expressions or Firedrake Functions onto a MixedFunctionSpace. We don't support assemble(interpolate(TestFunction(V), W)) yet.

2. We implemented restricted_function_space that returns an equivalent of RestrictedFunctionSpace constructed from a MixedFunctionSpace and a list of DirichletBCs to infer the individual boundary_sets of each subspace.

The PR includes basic tests for these new features.

[github-actions]

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake complex	8114 ran	6641 passed	1473 skipped	0 failed

[connorjward]

@pbrubeck FYI I think I may have accidentally broken the CI runners. I am investigating.

[github-actions]

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake real	8120 ran	7445 passed	675 skipped	0 failed

[pbrubeck]

@ksagiyam I'm not sure if importing DirichletBC on this file is a good idea

[ksagiyam]

DirichletBCs are big objects, so it would be nicer if impl.RestrictedFunctionSpace does not need to handle them, I think.
Though I said it would be good to have RestrictedFunctionSpace(mixed_space, ...) interface, I was not thinking too hard. Given that RestrictedFunctionSpaces are almost the same as regular FunctionSpaces (they basically only have different DoF orderings), they should just be viewed as sisters of regular FunctionSpaces, and we should probably just allow for MixedFunctionSpace([V_regular, V_restrict, V_regular, ...]) interface.
So, instead of changing the impl.RestrictedFunctionSpace constructor, can we simply convert the MixedFunctionSpace object all composed of regular function spaces to one composed of regular and restricted spaces as necessary inspecting bc.function_space in NonlinearVariationalProblem around here https://github.com/firedrakeproject/firedrake/blob/ed826a575a9314b45e8e396d1a2e9da8687f524d/firedrake/variational_solver.py#L91C45-L91C46?

[pbrubeck]

I'm not changing __init__, but just __new__. I think I achieved what you just described but through RestrictedFunctionSpace(MixedFunctionSpace, boundary_set=bcs). The bcs do not live within the RestrictedFunctionSpace, it is just easier to pass a single list as opposed to having to create separate boundary_sets for each mixed component.

[ksagiyam]

I suggest that we do not allow for RestrictedFunctionSpace(MixedFunctionSpace, ...) interface for the reason mentioned in the above. Passing bcs does not sound right either, no matter how we use it, as impl.RestrictedFunctionSpace is topological, while bcs contain geometric information.

[pbrubeck]

I think we could follow the same approach we have for FunctionSpace, and implement a user-facing function RestrictedFunctionSpace in functionspace.py that parses the bcs into a boundary_set for each component, and constructs the MixedSpace if a MixedSpace is provided.

[ksagiyam]

So why do you think we should allow for RestrictedFunctionSpace(MixedFunctionSpace, ...) interface?

[pbrubeck]

It would be the most natural and intuitive way to restrict a space, and would just work in general.

[pbrubeck]

@ksagiyam is this a reasonable place to include this code?

[pbrubeck]

@ksagiyam So I just rename this method to restrict_function_space(function_space, bcs) and then use it in variational_solver.py and eigensolver.py?

[connorjward]

Why not have a function_space.restrict(bcs) method? You are not being clear about the difference between functions and methods which is a little confusing.

[pbrubeck]

That is one idea I suggested during the meeting. I think that's the most intuitive to use. If implemented correctly, it even allows to further restrict an alredy restricted space.

[pbrubeck]

The class method function_space.restrict(bcs) is also really hard to implement, as there's a lot of confusion surronding the different topological and geometric classes...

[ksagiyam]

I think the biggest problem again is that we are still passing bc objects to make a restricted function space. Whatever makes restricted function space does not need to see the entire bc, it only needs bc.subdomain.

It is right that variational_solver.py and eigensolver.py want to share this part of the codes, and I think Connor is right in that we should do this with methods. So I think we could probably do something like the following:

mesh.py:

class SubdomainID:

    def __init__(self, meshes, ids):
        assert len(meshes) == len(ids)
        meshes = tuple(m.topology for m in meshes)
        self._meshes = meshes
        self._ids = ids

    def __len__(self):
        return len(self._ids)

    def __iter__(self):
        return iter(self._ids)

bc.py:

class DirichletBC:
    ...
    @staticmethod
    def extract_subdomain_id(bcs):
        ...
        return SubdomainID(..., ...)

functionspaceimpl:

class WithGeometry:
    ...
    def restrict(self, subdomain_id):
        return type(self)(self.topological.restrict(subdomain_id), self.mesh())

class MixedFunctionSpace:
    ...
    def reconstruct(self, subdomain_id):
        return type(self)(
            [V.reconstruct(sid) for V, sid in zip(self, subdomain_id)]
        )

class FunctionSpace:
    ...
    def restrict(self, subdomain_id):
        if subdomain_id is None:
            return self  # collapse?
        else:
            return RestrictedFunctionSpace(self, boundary_set=subdomain_id)

variational_solver.py:

class NonlinearVariationalProblem:
    ...
    def __init__(...):
        ...
        V_res = V.restrict(DirichletBC.extract_subdomain_id(bcs))
        ...

[connorjward]

By all means a RestrictedFunctionSpace doesn't need to know about bcs, but that doesn't mean we need to make a restrict method less ergonomic. Why can't we simply have something like:

def restrict(self, bcs_or_sid):
  if isinstance(bcs_or_sid, DirichletBC):
    sid = bcs_or_sid.subdomain_id
  else:
    sid = bcs_or_sid

  return RestrictedFunctionSpace(self, sid)

[ksagiyam]

The scope of each function/method should be as limited as possible; FunctionSpace should be "unable" to access bc.function_arg. This is in line with the concept of encapsulation (https://en.wikipedia.org/wiki/Encapsulation_(computer_programming)), I think.

[connorjward]

When it comes to object design I absolutely agree that we should be careful about this sort of thing. But I think with interface design we should prioritise ergonomics. Users are more used to dealing with function spaces and boundary conditions than SubdomainIds.

[ksagiyam]

In general, I think, if we design things correctly, we can come up with better/cleaner/more correct interface. I do not know where in Firedrake conceptual correctness is sacrificed for better user interface. SubdomainID could be used by users, but I intended that to be used in NonlinearVariationalProblem and EigenProblem; users can already construct mixed restricted space normally as V * V_res. The proposed interface, V.reconstruct(bcs), can not be used by users naturally, either, as bcs will have to be defined before making the space, and bcs must be reconstructed afterwards:

V = V1 * V2 * V3
bc = DirichletBC(V.sub(0), g, subdomain)
V_res = V.restrict([bc])
u = Function(V_res)
v = TestFunction(V_res)
F = inner(u, v) * dx
bc_res = bc.reconstruct(V=V_res)
sol = Function(V_res)
solve(F == 0, sol, bcs=[bc_res])

Given that bc.reconstruct simply makes a new object, this means that the original bc object was constructed merely to make V_res, while we know we only needed bc.subdomain_id; g should also be directly constructed on V_res. Note by the way that some users may want to work with RestrictedFunctionSpace directly, but I think the primal usage is solve(..., restrict=True), where users do not need to touch RestrictedFunctionSpace directly,