Integrating Firedrake and Deep Learning software

[jcddafonseca]

Hi!
I'm involved in a research project where we have to integrate Deep Learning models to Finite-Element solvers.
For many reasons, our team decided to use Firedrake, and the immediate technical problem that we face now is:
How can we make Deep-Learning software components, usually living in complex (and conda) environments, to communicate with Firedrake in its own virtual environment? There seems to be a few possible ways to solve this... Would the trivial one be bringing everything into Firedrake's environment (e.g., installing tensorflow, pytorch etc... or other machine-learning software), since we cannot take the opposite direction? My feeling is that this might work in some situations but it wouldn't be a good decision...

I'd like to hear from Firedrake community about the above questions: what do you think?

Thanks!
Julio Fonseca

[dham]

The ability to embed Firedrake in PyTorch is there right now. Passing --pytorch to the installer will install PyTorch. The coupling interface is described here: https://arxiv.org/pdf/2303.06871.pdf. If you want to extend this to other ML frameworks, then the code for FiredrakeTorchOperator is the pattern to follow.

On the right branches, the reverse coupling (PyTorch operators inside Firedrake) also works. @nbouziani can advise on what to use.

[nbouziani]

As @dham said, you can already embed Firedrake into PyTorch, using torch_operator (which relies on the FiredrakeTorchOperator interface). Embedding PyTorch into Firedrake, e.g. to include an ML-based model in a PDE system, is also possible on this branch, using neuralnet (which in turn relies on the PytorchOperator interface). For this coupling (PyTorch embedded in Firedrake), the interface is described in this second paper: https://arxiv.org/pdf/2111.00945.pdf.

I am also happy to help if you are interested in extending this to other ML frameworks.