UniqueDomainExtractor DAG traverser

[leo-collins]

Extract unique domain from an expression by traversing the expression tree. If we reach a terminal, we return the domain of that. If we reach an operator, we return the domain of the operands, if they all match.

Unlike the previous unique_domain_extractor, the DAG traverser works with expressions containing indexed arguments and coefficients. For example,

domain = MeshSequence([mesh1, mesh2])
V = FunctionSpace(domain, mixed_element)
u1, u2 = split(TrialFunction(V))
assert extract_unique_domain(u1) == mesh1
assert extract_unique_domain(u2) == mesh2

I plan to extend this in the future so that it works with BaseForms containing indexed arguments and coefficients, though I think this will require implementing an IndexedBaseForm class

[connorjward]

I wonder if this could be generalised a little to a "domain extractor" that returns all domains. extract_unique_domain is then a special case of that.

[ksagiyam]

This assumes that o.ufl_operands[0] has ufl_element(), so something like the following will fail:

cell = triangle
elem0 = LagrangeElement(cell, 1)
elem1 = LagrangeElement(cell, 1)
elem = MixedElement([elem0, elem1])
mesh1 = Mesh(LagrangeElement(cell, 1, (2,)), ufl_id=100)
mesh2 = Mesh(LagrangeElement(cell, 1, (2,)), ufl_id=101)
domain = MeshSequence([mesh1, mesh2])
V = FunctionSpace(domain, elem)
f = Coefficient(V)
d = extract_unique_domain(
    Indexed(
        grad(f),
        MultiIndex((FixedIndex(0), FixedIndex(1))),
    )
)

It's probably hard to make postorders robustly handle this sort of things.
First alternative that comes to my mind is preorder (or whatever you call it):

@process.register(Indexed)
def _(self, o, multiindex):
    expression, multiindex = o.ufl_operands
    return self(expression, multiindex=multiindex)

@process.register(Grad)
def _(self, o, multiindex):
    return self(o.ufl_operands[0], multiindex[:-1])

... and so on

You can see https://github.com/FEniCS/ufl/blob/main/ufl/algorithms/apply_coefficient_split.py.
CoefficientSplitter works as we make sure that ReferenceValue, ReferenceGrad, and Restricted have already been propagated to directly wrap terminals by the time we use it; you might need some "balancing" of this kind before calling extract_unique_domain, too, but it could be inconvenient.

Do you only need this for indexed Arguments and Coefficients? What is the most general use case in your mind?

[ksagiyam]

Should define separate singledispatch method as:

@process.register(Coefficient)
@process.register(Argument)
def _(self, o, ...):
    ...

Same for the other types.

[ksagiyam]

Is this what we expect?

[ksagiyam]

Can we not change the definition of extract_unique_domain() here? Instead, can we change

        for t in traverse_unique_terminals(expr):
            domainlist.extend(t.ufl_domains())

in extract_domains() function to use DAGTraverser? I think it is important to make sure that extract_domains() and extract_unique_domains() are consistent.