Force Ad operators to have domains (and ranges)

[keileg]

In connection with work on #1497, I need to make decisions based on the domains of Ad operators, that is, the subdomains, interface grids, or boundary grids on which they are defined. I can give details if someone is curious, but it is not important in this context. Though I can hack my way through the issues I am facing, it is more natural to require that all operators have a domain.

To be clear, I am only considering domains in the form of subdomains and interfaces, and not the finer grained view of, say, operators defined on cells or faces on a grid. We may want to consider that as well (I have a half-baked idea in mind, also motivated by faster Ad parsing), but that will have to come later.

What can be gained:

1. The project Implementation of a dense array backend for forward AD #1497, which is expected to give considerable computational savings in some settings, can be implemented without dirty tricks.
2. It becomes possible to do consistency checks so that, for instance, we cannot add operators defined on two different domains. I believe we can grow this idea further to ease debugging of Ad expressions considerably (though I will prefer to consider that a separate project).

What concerns me:

1. An operator must be assigned the correct domain whenever it is defined. Though it should always be clear what that domain is (if this is not the case, it is good that the developer is forced to stop and think, because then something is wrong), assigning the domains to the operator will feel like writing boilerplate code. I would argue this will be a limited problem after the initial roll-out.

Technical details I have hammered out so far:

1. Scalars are a special case, these have no domains.
2. Mortar and subdomain projections can have both a domain and a range. Hence, all operators should have both, with the default option that they are equal. Exactly how to implement this is not clear, but it should not be a difficult decision.

Does anyone have thoughts or suggestions?

[IvarStefansson]

Have you started thinking of vector quantities?
To your technical details:

1. We could also consider them to formally have domains and ranges, with implicit broadcasting.
2. I think all operators have ranges. The distinction for projections is that the two may differ.

[keileg]

Vector quantities should make no difference as long as we only consider grids, hence no counting of actual dogs. One step at a time.

I agree, they should all have explicit range and domain.

I don't know what you mean by broadcasting in this context, let's discuss in person.

[keileg]

Note to self: Scalars should either have no domain or some special domain value which signifies it can be combined with anything.

[keileg]

The design is getting clearer to me now. Though the motivation in the original post was for operators to have a domain for certain checks that are needed in #1497, we can do better: With a clean implementation, it should not take much code to check the validity of an Ad operation at the time of construction - right now, we often get a cryptic error message at the time of parsing instead.

The envisioned implementation steps are as follows:

1. The GridEntity type in equation_system.py is converted to an enum and used wherever we now write "cells" etc. For now, I have only considered use in the context of dof-management, broadly speaking, wider use has not been thought of, but if relevant, this can be done later.
2. The dof_info attribute that some operators have, will be represented (note: not an attribute) in all operators. I would prefer to represent the dof_info as a dataclass that maps GridEntity to ints, as this seems more principled, but the choice of data structures is secondary for now.
3. We introduce an enum for representing the type of a domain an operator is associated with (more below). Right now, this information is carried in a string attribute that can (by typing) take values subdomains, interfaces, boundary grids. A fourth type void, representing ad scalars should be added. I don't have a good name for the enum, or the type of information it carries, but will refer to it as domain type below.
4. We introduce a 'collection' (could be a namedtuple, but we need to think carefully of how equality of collections is evaluated) of the triplet {domain type, grids, dof_info}, where 'grids' is what is now represented in pp.ad.Operator.domain, that is, a list (I think tuple would be better, but that is a different matter) of subdomains, interfaces, or boundary grids. Although the 'domain type' property can be inferred from the 'grids', it will likely be beneficial to represent it explicitly - after all, there is a reason our current operators have a domain_type attribute. A good name is needed for this collection.
5. When constructing an atomic Ad Operator (that is, a leaf in the operator tree, not a composed operator), its domain and range should be explicitly specified as two collections specified in point 4. For most operators, the required information should be readily available, but it could be that we want some convenience methods for easy construction (this I have not thought about). The exception that I have found is spatial discretizations, represented as MergedOperators: These must get hold of the dof_info from the discretization class, hence the latter needs a get-method for that information. This should be a minor change.
6. The Operator enum, residing in pp.ad.operators.py, gets a method validate_operands. For most operations, this verifies that the operands have the same domains and the same ranges (that is, all items in the point 4 collection are the same). The exception is matmul ('@'), which compares the domain of the first operand with the range of the second. Failure of the validation raises an error.
7. Assuming the operands are compatible, the combined operator has its domain and range inferred from the operands, likely as a method infer_domain_range in the Operator enum. Again matmul will be a special case.

Thoughts on testing:

• Equality of the collections (point 4) should be thoroughly tested.
• If we implement the convenience method for collection construction (point 5), it can be tested.
• validate_operands need thorough testing.
• infer_domain_range should be tested.
• Ultimately, the testing of models will give excellent integration (system?) testing.

What can fail/assumptions:

• The integrity of the validate_operands method may deteriorate for operands with more than one dof per grid entity, or with dofs for multiple grid entities (say, cell and face). The check outlined above cannot safeguard against numbering incompatibilities - is the cell index or the variable component index the fast one? This is unfortunate, but this requires a different project, I think.

Future extensions:

• With the above system in place, it should be straightforward to validate operator units (in the SI sense) and propagate this information for combined operators. It will however require that we define units for all operator leaves, and possibly that we have a better way of representing units (or pp.Units could be excellent for this purpose, I have not considered it at all). This should however be a separate project.