Define custom UFL object (UserExpression) with Neural Network

Hi all,

I want to implement a custom UFL object that can replace a part of my variational formulation with a neural network (PyTorch). The following picture should explain the concept a little bit better:

diagram.drawio
(yellow objects represent UFL objects, the blue object is the UserExpression which transforms its inputs according to the networks architecture but “looses” the information about the input u)

This way, I loose information about my input parameter (in this case the TrialFunction) to the Network, as the UserExpression does not seem to carry these information with it. All in all, FEniCS is not able to solve the problem, since a (the left hand side) is not dependent on u anymore:

*** Error:   Unable to define linear variational problem a(u, v) == L(v) for all v.
*** Reason:  Expecting the left-hand side to be a bilinear form (not rank 1).
*** Where:   This error was encountered inside LinearVariationalProblem.cpp.

In addition to that, it would not be able to solve the (nonlinear) problem e.g. with the NewtonSolver, as this would require a pure implementation using UFL objects (as stated for example here or here). The iterations would have to be implemented outside of FEniCS.

What would be required to have the UserExpression be a “first class” UFL-object that is fully compatible with the assembly of the variational formulation? Some ressources or entry-points for the development would be very helpful!

Thanks in advance!

Hi, a possible remedy is to define your \sigma (if I understand this is your NN output) as some Function on a suitable FunctionSpace and update its values coming from the NN within each iteration of a custom Newton solver. You could do the same with the tangent operator needed for the Jacobian matrix. A similar strategy is used here for implementing plasticity problems