Programmable equation transformations

It transpires that people would like to be able to write (weak) equations, and then apply transformations to them before using them in Firedrake variational solvers. Examples of such transforms include:

• Applying SUPG to some or all terms.
• Applying one of several time stepping schemes.

These transformations conceive of a an equation as a sum of terms, and the transformations are applied selectively to terms. There are therefore two problems. First, deciding which transformations are applicable to each term, and second, applying those transformations. The second is actually a solved problem in the sense that the transformations involved can all be encoded by UFL multifunctions, and indeed mostly by using replace. The first requires a little thought.

Let us define an Equation class which will be a sum of terms. Each of these terms is a form (noting that forms are closed over addition) which carries one or more labels. A label is an object containing key, value pair where value defaults to True. This enables labels which are essentially independent flags, but also mutually exclusive sets (same key, different values). The latter facilitates support for mutually incompatible options such as the way a term should be mapped to the time stepping scheme.

Labels and equations also need to support various operations which will enable reasoning about the equations by setting labels based on existing label values. Consider the following:

mass = Label("mass")
advection = Label("advection")
diffusion = Label("diffusion")
m = <mass form>
a = <adf_form>
d = <diff_form>
eq = mass(m) + advection(a) + diffusion(d)

These operations are defined as follows:

1. Calling a Label on a Form produces a Term encoding that form and that label.
2. Calling a Label on a Term produces a new Term with the label added to the existing set.
3. Calling a Label on an Equation produces a new Equation with the label added to each Term.
4. The sum of two Terms or a Term and an Equation is an Equation with the corresponding terms.
5. Label should also have a remove method such that calling l.remove(t) returns a new Term with label l removed.

Note in every case that operations return new objects. It is an important rule of computer symbolic algebra that you never change symbolic objects.

Note also that users never need to instantiate Term or Equation objects. They just label forms and add them together.

Label transformations.

The point of labelling terms is to enable reasoning to eventually determine which terms to apply. Let's take a simple example which labels terms for a time stepper:

implicit = Label("time", "implicit")
theta = Label("time", "theta")
eq = eq.label_map(advection, theta)
eq = eq.label_map(diffusion, implicit))

This is adds the right hand label to each term for which the left hand is true. In order to make this more flexible, we should define the logical operators over Label objects. This is easily done with the appropriate Python magic methods. These should return symbolic expressions which we can evaluate by substituting the current values of the labels on each Term. This would enable, for example:

eq = eq.label_map((label_a and not label_b, label_c.remove))

LM comments. Unfortunately, it is not possible to provide magic methods to override and, or, or not. You can only provide __bool__ which coerces an object to bool. We could instead overload the bitwise operators so you could write this as label_a & ~label_b, but that's quite ugly.

Because these are straight Python expressions, it's also possible to put any Python expression you like in the left hand side and its truth value will be used. The general form of label_map is:

eq.label_map(filter, map_function)

Where filter is an expression comprising general Python expressions and Labels, linked by Boolean operations. map_function is a function which takes in a Term and returns a Term, an Equation, or None. Returning None indicates that the term concerned is to be dropped. eq,label_map returns a new Equation whose terms are the result of applying map_function to those Terms of eq for which filter is true. Where filter is false, the terms of eq are carried through to the new Equation.

It is claimed that this amount of programmability, or something very close to it, will enable arbitrarily complex relabelling.

Term transformations

The label_map method actually also provides the mechanism by which the actual term transformations can be applied. To see this, note that the right hand function in each key-value pair is simply a function which takes in a Term and returns a Term. It might be useful to enable it to return an Equation as well so that Terms can be split into multiple terms (e.g. an implicit and an explicit part).

Other considerations.

• Equation should have a .form member which strips the labels and adds the form in each Term. This enables the values to finally be passed to Firedrake.
• Some care will be needed to establish the core taxonomy of labels so that time steppers are composable with equations.
• It appears to be necessary to label the unknown in each term, because it's actually not possible to discover it by inspection.