High-order differentiation of forms with external operators

[a-latyshev]

There is an issue concerning the high-order differentiation of forms that contain ExternalOperator-s.

We start with a standard form and its first derivative

N = ExternalOperator(u, function_space=V_ufl) # N(u; v')
E = N * dx # N(u; v') * dx
F = derivative(E, u, u_test)
F_expanded = expand_derivatives(F) # Action( v1 * dx; dN(u; v1, v'))

Then we want to differentiate F or F_expanded, but unfortunately, there is an error in formoperators when we differentiate Action (see the error below)

# Differentiation expansion -> Second derivative -> Second expansion
J_expanded = derivative(F_expanded, u, u_trial)
J_expanded_expanded = expand_derivatives(J_expanded)  # Error

Without the F expansion, this works, but this is worse since one gets 0, which is not correct:

# Double differentiation
J = derivative(F, u, u_trial)
J_expanded = expand_derivatives(J) # == 0 <--- Works, but incorrectly

It seems that Action is not properly differentiated at some point. The AD algorithm cannot handle the situation where the right part has to be differentiated as well as the left one.

An obvious workaround is to manually define a new form F_impl, where we explicitly use the right part of F_expanded, which is the derivative of N. Basically, we replace the action F_expanded with a normal form:

# Replace Action with a form explicitly
dN = F_expanded.right() # dN(u; v1, v')
F_expl = dN * u_test * dx
J_expl = derivative(F_expl, u, u_trial)
J_expl_expanded = expand_derivatives(J_expl) # Action( v1 * v2 * dx; ddN(u; v2, v1, v')) <--- Works correctly

How should one handle the issue? Introduce a replacement mechanism for Action before differentiation, or is it possible to augment the current algorithm recursively?

Full MWE:

from ufl.core.external_operator import ExternalOperator
from ufl import Coefficient, Argument, derivative, dx
from ufl.algorithms import expand_derivatives
from mpi4py import MPI
from dolfinx.mesh import create_rectangle, CellType
from dolfinx.fem import functionspace

domain = create_rectangle(
    MPI.COMM_WORLD,
    [[0.0, 0.0], [1.0, 1.0]],
    [10, 10],
    cell_type=CellType.triangle,
)
V = functionspace(domain, ("Lagrange", 1))

V_ufl = V.ufl_function_space()
u = Coefficient(V_ufl)
u_test = Argument(V_ufl, 1)
u_trial = Argument(V_ufl, 2)
N = ExternalOperator(u, function_space=V_ufl) # N(u; v')
E = N * dx # N(u; v') * dx
F = derivative(E, u, u_test)
F_expanded = expand_derivatives(F) # Action( v1 * dx; dN(u; v1, v'))

# Replace Action with a form explicitly
dN = F_expanded.right() # dN(u; v1, v')
F_expl = dN * u_test * dx
J_expl = derivative(F_expl, u, u_trial)
J_expl_expanded = expand_derivatives(J_expl) # Action( v1 * v2 * dx; ddN(u; v2, v1, v')) <--- Works correctly

# Double differentiation
J = derivative(F, u, u_trial)
J_expanded = expand_derivatives(J) # == 0 <--- Works, but incorrectly

# Differentiation expansion -> Second derivative -> Second expansion
J_expanded = derivative(F_expanded, u, u_trial)
J_expanded_expanded = expand_derivatives(J_expanded)  # Error

Error:

---> [36](vscode-notebook-cell:?execution_count=18&line=36) J_expanded = derivative(F_expanded, u, u_trial)
     37 J_expanded_expanded = expand_derivatives(J_expanded)  # Error

File /dolfinx-env/lib/python3.12/site-packages/ufl/formoperators.py:411, in derivative(form, coefficient, argument, coefficient_derivatives)
    409     dright = derivative(right, coefficient, argument, coefficient_derivatives)
    410     # Leibniz formula
--> [411](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a2266656e696373782d302e31302e30227d.vscode-resource.vscode-cdn.net/dolfinx-env/lib/python3.12/site-packages/ufl/formoperators.py:411)     return action(
    412         adjoint(dleft, derivatives_expanded=True), right, derivatives_expanded=True
    413     ) + action(left, dright, derivatives_expanded=True)
    414 else:
    415     raise NotImplementedError(
    416         "Action derivative not supported when the left argument is not a 1-form."
    417     )

File /dolfinx-env/lib/python3.12/site-packages/ufl/formoperators.py:205, in action(form, coefficient, derivatives_expanded)
    203     if isinstance(form, Form):
    204         return compute_form_action(form, coefficient)
--> [205](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a2266656e696373782d302e31302e30227d.vscode-resource.vscode-cdn.net/dolfinx-env/lib/python3.12/site-packages/ufl/formoperators.py:205) return Action(form, coefficient)

File /dolfinx-env/lib/python3.12/site-packages/ufl/action.py:91, in Action.__new__(cls, *args, **kw)
...
--> [217](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a2266656e696373782d302e31302e30227d.vscode-resource.vscode-cdn.net/dolfinx-env/lib/python3.12/site-packages/ufl/action.py:217)     V_left = left.arguments()[-1].ufl_function_space().dual()
    218 else:
    219     raise TypeError("Action left argument must be either Coefficient or BaseForm")

IndexError: tuple index out of range

[jorgensd]

Here is the same code, with a pure ufl implementation (has to be executed from ufl/test to get the helpers for dummy finite elements:

import ufl
from ufl.core.external_operator import ExternalOperator
from ufl import Coefficient, Argument, derivative, dx
from ufl.algorithms import expand_derivatives

#
from ufl.pullback import identity_pullback
from ufl.sobolevspace import H1

from utils import FiniteElement, LagrangeElement

domain = ufl.Mesh(LagrangeElement(ufl.triangle, 1, (2,)))


f1 = FiniteElement("Lagrange", ufl.triangle, 1, (), identity_pullback, H1)

V = ufl.FunctionSpace(domain, f1)

u = Coefficient(V)
u_test = Argument(V, 1)
u_trial = Argument(V, 2)
N = ExternalOperator(u, function_space=V)  # N(u; v')
E = N * dx  # N(u; v') * dx
F = derivative(E, u, u_test)
F_expanded = expand_derivatives(F)  # Action( v1 * dx; dN(u; v1, v'))

# Replace Action with a form explicitly
dN = F_expanded.right()  # dN(u; v1, v')
F_expl = dN * u_test * dx
J_expl = derivative(F_expl, u, u_trial)
J_expl_expanded = expand_derivatives(
    J_expl
)  # Action( v1 * v2 * dx; ddN(u; v2, v1, v')) <--- Works correctly

# Double differentiation
J = derivative(F, u, u_trial)
J_expanded = expand_derivatives(J)  # == 0 <--- Works, but incorrectly

# Differentiation expansion -> Second derivative -> Second expansion
J_expanded = derivative(F_expanded, u, u_trial)
J_expanded_expanded = expand_derivatives(J_expanded)  # Error

[jhale]

@nbouziani Any ideas on this issue? We tried debugging it but couldn't see the solution. We can finish up a proper PR with tests if you can point us in the right direction on a fix.

[dham]

I haven't got time to look into the details of this right now (it's the crazy part of term) but the error message gives me a hint. Basically there are cases where naively applying the chain rule will produce an Argument where you ought to have a form. This happens because of reflexivity (V** being naturally isomorphic to V). This means that an Argument is equivalent to the identity form in V x V*->R and that mixup can sometimes happen in the chain rule. I think there are already some sui generis fixes in the derivative code for this case, but clearly not enough.

I would be walking through the differentiation process in a debugger looking for the point at which you get an Expression (which I'm pretty sure will be an Argument) dropping into an Action.