Function to extract linear combination if present in ufl_expression

[jorgensd]

Typically, one would like to assign a linear combination of functions in the same function space to each other (for instance during time-stepping).
I encountered this when trying to get a DOLFINx Irksome implementation up and running:
firedrakeproject/Irksome#234
and in Legacy DOLFIN there is a rather large complicated section on this extraction:
https://bitbucket.org/fenics-project/dolfin/src/a8dcf753d3ab9ae19d7bedf8c634a4fe545032d2/python/dolfin/function/function.py?at=master#lines-37:177

Instead of repeating the same logic over and over in different implementations (legacy Dolfin, DOLFINx, firedrake) it would be nice to have this function native in UFL.

[michalhabera]

See also https://github.com/fenics-dolfiny/dolfiny/blob/d2456e7d64c6de15384e219fbd0f8b1b0a81b544/src/dolfiny/expression.py#L143-L201 for a simple version of this. Useful also in a time-stepping context.

[connorjward]

For a less simple version see https://github.com/firedrakeproject/firedrake/blob/main/firedrake/assign.py#L35. Something in UFL would be nice.

[jorgensd]

So I guess what we want is a combination of the two, preferably using the DagTraverser?

[connorjward]

So I guess what we want is a combination of the two, preferably using the DagTraverser?

Yeah I think so. Note that our implementation is made more complicated by the existence of the Real space. In Firedrake a Function in R is the same thing as your UFL Constant. So in Firedrake we allow the user to treat a Function in R as the scalar in a linear combination. Ideally the implementation here will be extensible (via another @process.register?) so we can add a custom handler for that inside Firedrake.

We also want to allow the multiplication of vector-valued functions by vector-valued constants.

[jorgensd]

So I guess what we want is a combination of the two, preferably using the DagTraverser?

Yeah I think so. Note that our implementation is made more complicated by the existence of the Real space. In Firedrake a Function in R is the same thing as your UFL Constant. So in Firedrake we allow the user to treat a Function in R as the scalar in a linear combination. Ideally the implementation here will be extensible (via another @process.register?) so we can add a custom handler for that inside Firedrake.

We also want to allow the multiplication of vector-valued functions by vector-valued constants.

I have a similar thing for real functions in the Irksome-dolfinx PR.

So I guess what we want is a combination of the two, preferably using the DagTraverser?

Yeah I think so. Note that our implementation is made more complicated by the existence of the Real space. In Firedrake a Function in R is the same thing as your UFL Constant. So in Firedrake we allow the user to treat a Function in R as the scalar in a linear combination. Ideally the implementation here will be extensible (via another @process.register?) so we can add a custom handler for that inside Firedrake.

We also want to allow the multiplication of vector-valued functions by vector-valued constants.

So if I have a field w=(u, v) where each component is separate (not Nedelec/RT etc.) you would like to say take C=(a, b) and with ufl.dot(c, w) extract [(u, a), (v, b)]?

If make make a DAGTraverser that handles the basic UFL cases (sum, product, division, Componenttensor etc) it should be easy to extend it to new classes by inheritance.

[connorjward]

So if I have a field w=(u, v) where each component is separate (not Nedelec/RT etc.) you would like to say take C=(a, b) and with ufl.dot(c, w) extract [(u, a), (v, b)]?

Actually I think I hallucinated that this works for us. If we have

from firedrake import *

mesh = UnitSquareMesh(3, 3)
V = VectorFunctionSpace(mesh, "P", 1)
c = Constant([3, 4])
f = Function(V)

then f.assign(c) is allowed but c*f crashes with an invalid rank error.