Extend derivative block to derive residuals from functionals

[Joseabcd]

Make fem.forms.derivative_block work on functionals to derive residuals, as described in #3766.

The following notes could be useful for the review:

1. The new logic is in a new code paragraph, independent of the original code. Not sure if a different style is preferred
2. I have used ValueError exceptions as opposed to asserts, assuming that the checks are on user input (as opposed to checks against bugs). The original code however has some asserts, so I’m not sure which is their right intent
3. I have not covered the case where F is a list of functionals
4. The original code creates a list of trial functions independently. Perhaps it should be updated to use ufl.TrialFunctions(ufl.MixedFunctionSpace(…))?
5. As in the original code, I have not checked that du is the correct type of argument (test vs trial), but possibly this might be helpful

[schnellerhase]

Looks like a good start.

• Your case switching check is identical for rank 0 to 1 and 1 to 2. A dependency on the rank will be necessary to differentiate. For a ufl form a you can access the rank as a.rank. I now saw the arguments check.
• To ensure your code works and to simplify development unit tests would be good to add, for example in test_forms.py. (The rank 1 to 2 case seems to be only unit tested through the assembler at the moment, would be good to also add a small unit test for this case)

[Joseabcd]

thanks @schnellerhase and @jorgensd

I found no rank() method, so I used the more cumbersome check with len(F.arguments())

While adding tests, I realized there were more situations where the user can supply bad inputs and get obscure errors (or unintended results silently), so I have been more comprehensive when checking the inputs.

For example, before you could wrongly pass trial functions to derive functionals, and it would still return a form (or Nones) without complaining, asufl.derivative seems quite silent. I’ve also made sure that incorrect arguments don't exit derivative_block with an unhelpful error. After this,derivative_block grew, so its body is now grouped in small functions to make it more readable.

For consistency, I’ve also replaced the asserts with ValueError. And I’ve used ValueError even in cases when the strictly correct exception should be TypeError, in order to avoid more extra conditions (assuming that users won’t care and that they won't be dealing with these exceptions at such level of detail). I was not quite sure about the preference on this point though

Could you let me know how it looks, if the new checks do indeed belong here, or any other feedback or thoughts?

[Joseabcd]

In this case, where u and du are not sequences (what I called "univariate" functional), there is a subtle choice I made worth mentioning.

In the odd event that the user passes v0 from a mixed space, applying ufl.extract_blocks to the output of ufl.derivative will be a sequence containing one ufl.Form. In other words, the result can be a ufl.Form or a Sequence depending on whether v0 is part or not of a mixed space (!)

I found this potentially confusing, so I avoided applying ufl.extract_blocks in derivative_block when the input du is just a function

[schnellerhase]

This should be added to the documentation.

[Joseabcd]

@schnellerhase I'll update the docstring, but could you tell me if I made the right choice? I'm still exploring the rest of the project and I lack the big picture of how this fits in the rest of the python layer.

I think that in essence, my dilemma is captured in the following question:

Given:

mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
V0 = functionspace(mesh, ("Lagrange", 1))
V1 = functionspace(mesh, ("Lagrange", 2))
V = MixedFunctionSpace(V0, V1)
f0, f1 = dolfinx.fem.Function(V0), dolfinx.fem.Function(V1)

M = f0**2 * dx

which of the following is the more accurate statement?

• M is a functional defined over a single function. We don't know about "blocks".
• M is a functional defined over N functions (where N happens to be 1). We work with "blocks".

[schnellerhase]

Going along the argument in #4118 (comment) - I would say the first option is the favourable perspective. Going with the perspective: the block interpretation is at the symbolic level one of argument counting, not one of the setup of the argument's function spaces. What do you think @jorgensd - does that make any sense?

[Joseabcd]

This should be added to the documentation.

I implemented in the docstring a more structured format that I observed in other functions from the project. I hope the detail about the type of returned output depending on the input is clearer now

[jorgensd]

I also think the first one is the correct thing to consider here. As f1 does not enter the functional, the construct of blocking them is of no interest to the functional itself.

[jhale]

Shouldn't this be raised within ufl.derivative already?

[jorgensd]

It is not raised, ref:

from ufl import (
    FunctionSpace,
    Mesh,
    Coefficient,
    dx,
    TrialFunction,
    derivative
)
import basix.ufl
from mpi4py import MPI
import dolfinx
cell = "interval"
P1 = basix.ufl.element("Lagrange", cell, 1, shape=())
domain = Mesh(basix.ufl.element("Lagrange", cell, 1, shape=(1,)))
V = FunctionSpace(domain, P1)

u = Coefficient(V)
J = u ** 2 * dx
v = TrialFunction(V)
F = derivative(J, u, v)
F_compiled = dolfinx.fem.compile_form(MPI.COMM_WORLD, F)

which runs with no problem, but becomes a problem if you want to get the jacobian

[Joseabcd]

Should we then keep/remove this check? In my view, if the concept of deriving a rank-zero form w.r.t. a “trial” function never ever makes sense (that is, if it’s conceptually wrong at the UFL layer), then it’d be UFL’s responsibility to complain

[jorgensd]

ufl strictly speaking actually handles this correctly.
If we consider this minimal example using the Lagrange-element from the UFL test suite:

from ufl import (
    FunctionSpace,
    Mesh,
    Coefficient,
    dx,
    TrialFunction,
    derivative,
    interval,
    TestFunction
)

from utils import LagrangeElement

cell = interval
P1 = LagrangeElement(cell, 1, shape=())
domain = Mesh(LagrangeElement(cell, 1, shape=(1,)))
V = FunctionSpace(domain, P1)

u = Coefficient(V)
J = u ** 2 * dx
v = TestFunction(V)
F = derivative(J, u, v)
N = derivative(F, u)
for arg in F.arguments():
    print("Test F", arg.number(), arg.part())
for arg in N.arguments():
    print("Test N", arg.number(), arg.part())

v = TrialFunction(V)
F_t = derivative(J, u, v)
N_t = derivative(F_t, u)
for arg in F_t.arguments():
    print("Trial F", arg.number(), arg.part())
for arg in N_t.arguments():
    print("Trial N", arg.number(), arg.part())

this returns:

Test F 0 None
Test N 0 None
Test N 1 None
Trial F 1 None
Trial N 1 None
Trial N 2 None

which means that if you do not supply the derivative direction of your second variation, you get an argument that has its number incremented by one, which is then correct in UFL terms, as the first variation is the one with the lowest numbering.

[Joseabcd]

@jorgensd was your comment in favor of removing or keeping the check? I was waiting for others to comment too.

I see UFL’s approach is to simply apply the math operation requested by the user (the variational problem is not UFL layer’s business). The question that remains is whether it is derivative_block’s business. If yes, then it should check that test/trial arguments are correct.

I am myself not sure. On one hand it would be convenient if derivative_block complains about a poorly defined problem. On the other hand, removing the check gives the user more flexibility to define the problem in unconventional ways; for example, deriving a functional first w.r.t. a trial function, and then w.r.t. a test function (which results in a seemingly-valid rank-2 form, with argument numbers 0 and 1).

@jhale what do you think? I am ok to proceed either way. Could you let me know any concerns about this and the rest of changes? Happy to work on them

[jorgensd]

I think checking it becomes complicated, as seen by the example above, where number!=1 if one uses derivative on a linear form (which has been defined with a trial function).

So i am in favor of removing the check and leave it to the user to define their forms correctly.

[Joseabcd]

I have removed the check now

[jhale]

Shouldn't this be raised within ufl.derivative or ufl.extract_blocks already?

[Joseabcd]

The UFL functions aren’t raising here either. For example, ufl.derivative(F, [u0, u1, u2], [du0]) silently returns a sequence with just one form. But I agree this is UFL’s job, so I have for now removed the check len(u) == len(du).

I will wait to remove the other checks that remain in this if, depending on the answer to question in the previous comment.

[Joseabcd]

Following the answer in linked comment, I’ve removed now the check that remained

[jhale]

Too picky, let UFL handle it.

[Joseabcd]

Before removing it, we need to answer the question in the comment above: should we allow derivative_block to derive a functional using a “trial” function?

[Joseabcd]

Following the answer in linked comment, I have removed it now

[jhale]

Let UFL handle it.

[Joseabcd]

The tests isinstance(du, Sequence)and len(u) == len(du) I believe are appropriate. Given that the concept of “component-wise” Jacobian is created here in the Python layer, isn't this layer responsible for making sure the result is a square “matrix”? UFL will never get to see anything other than single pairs of (u_i, du_i), so will never get a chance to complain.

About the third check on du_i.number(), I will wait to remove it depending on answer to question in this comment above

[Joseabcd]

Following the answer in linked comment, I’ve removed now the du_i.number() check

[jhale]

What does not F.arguments() do/why does it work?

[jorgensd]

F.arguments() returns a tuple of ufl.Argument. if a tuple is empty, its boolean state is False, thus the check says, if we have no arguments (rank=0) go there.

[jhale]

This looks pretty good but I think it's a bit 'strongly typed/overly checked' with if statements - shouldn't a lot of these already be picked up lower down in UFL?

[schnellerhase]

I found no rank() method, so I used the more cumbersome check with len(F.arguments())

My bad - the rank functionality is only available for expressions. Should be added to the UFL maybe as form.arity property. len(F.arguments()) is good here.

For example, before you could wrongly pass trial functions to derive functionals, and it would still return a form (or Nones) without complaining, asufl.derivative seems quite silent.

I agree this is a problematic interface. The underlying reason for this is that the only difference between a trial and test function in UFL is an index denoting the placement https://github.com/FEniCS/ufl/blob/61ea6299409a61d9e90e94d0b34f14ed269fa745/ufl/argument.py#L297-L312, however no input checks are performed on 2 only showing up if 1 is provided. This causes also other surprising behaviour #3721.

[Joseabcd]

This looks pretty good but I think it's a bit 'strongly typed/overly checked' with if statements - shouldn't a lot of these already be picked up lower down in UFL?

I've removed now some of the checks as suggested above. I'm happy to remove the rest of those suggested, but I left a question in an inline comment that I think needs to be resolved first

I marked as "resolved" the comments that I believe require no further action, but please let me know if anything doesn't look good

[Joseabcd]

I think all comments above have been addressed now. The stringent checks on the inputs have mostly been removed, letting inputs reach UFL in most cases. Some checks remain inside _derive_block_jacobian (see comment), to prevent from breaking the code that here calculates the component-wise Jacobian.

Could you please have another look at it when you have time, and let me know if there are still any concerns?