Escaping the abstraction: a foreign function interface for the Unified Form Language [UFL]

TL;DR

This paper introduces a foreign function interface for the Unified Form Language within the Firedrake finite element system, enabling seamless integration of deep learning models with PDE solvers for enhanced physical modeling.

Contribution

It presents a novel interface that connects Firedrake with deep learning libraries like PyTorch, supporting complex PDE and deep learning couplings while preserving separation of concerns.

Findings

Supports automatic differentiation with deep learning models

Enables automated inverse problem solving

Interfaces with PyTorch for flexible ML integration

Abstract

High level domain specific languages for the finite element method underpin high productivity programming environments for simulations based on partial differential equations (PDE) while employing automatic code generation to achieve high performance. However, a limitation of this approach is that it does not support operators that are not directly expressible in the vector calculus. This is critical in applications where PDEs are not enough to accurately describe the physical problem of interest. The use of deep learning techniques have become increasingly popular in filling this knowledge gap, for example to include features not represented in the differential equations, or closures for unresolved spatiotemporal scales. We introduce an interface within the Firedrake finite element system that enables a seamless interface with deep learning models. This new feature composes with the automatic differentiation capabilities of Firedrake, enabling the automated solution of inverse problems. Our implementation interfaces with PyTorch and can be extended to other machine learning libraries. The resulting framework supports complex models coupling PDEs and deep learning whilst maintaining separation of concerns between application scientists and software experts.