TensorDiffEq: Scalable Multi-GPU Forward and Inverse Solvers for Physics Informed Neural Networks

As part of the burgeoning field of scientific machine learning (Baker et al., 2019), physics-informed neural networks (PINNs) have emerged recently as an alternative to traditional partial different equation (PDE) solvers (Raissi et al., 2019; Raissi, 2018; Wight and Zhao, 2020; Wang et al., 2020b), and have given rise to the larger field of study in neural network approximation of PDE systems, generally referred to as Neural PDEs. Typical black-box deep learning methodologies do not take into account the underlying physics of the problem domain. The Neural PDE approach is based on constraining the output of a deep neural network to satisfy a physical model specified by a PDE. PINNs typically perform this task via PDE-constrained regularization of a residual function defined by the approximation of the solution network u and forward-pass calculations through the physics of the PDE model, with the applicable derivatives of u calculated via reverse-mode automatic differentiation in a modern deep learning framework such as Tensorflow (Abadi et al., 2016).

The potential of using neural networks as universal function approximators to solve PDEs had been recognized since the 1990’s (Dissanayake and Phan-Thien, 1994; Lagaris et al., 1998). However, Physics-Informed Neural Networks promise to take this approach to a different level through deep neural networks, the exploration of which is now possible due to the vast advances in computational capabilities and training algorithms since that time (Abadi et al., 2016; Revels et al., 2016) and modern congenial automatic differentiation software (Baydin et al., 2017; Paszke et al., 2017).

A great advantage of the PINN architecture over traditional time-stepping PDE solvers is that the entire spatial-temporal domain can be solved at once using collocation points distributed quasi-randomly (rather than on a grid) across the spatial-temporal domain, in a process that can be massively parallelized via GPU. As we have continued to see GPU capabilities increase in recent years, a method that relies on parallelism in training iterations could begin to emerge as the predominant approach in scientific computing. To this end, while other software suites exist to define and solve PINNs (Rackauckas and Nie, 2017; Lu et al., 2021; Hennigh et al., 2020; Haghighat and Juanes, 2021), many of those platforms are either restricted to single-GPU implementation or are not fully open-source. Additionally, with full support and customization capabilities of the Keras neural network ecosystem built in to the package, researchers and practitioners can define and train their own custom neural network architectures to approximate the solution of their problem domains. TensorDiffEq provides these scalable, modular, and customizable multi-GPU architectures and solvers in a fully open-source platform, tapping into the collective intelligence of the field to improve the implementation of the software and provide input on the direction, structure, and feature coverage of the framework.

Consider a general nonlinear PDE of the form:

	$\displaystyle\mathcal{N}_{\boldsymbol{x},t}[u(\boldsymbol{x},t)]=0\,,\ \ \boldsymbol{x}\in\Omega\,,\ t\in[0,T]\,,$		(1)
	$\displaystyle u(\boldsymbol{x},t)=g(\boldsymbol{x},t)\,,\ \ \boldsymbol{x}\in\partial\Omega,\ t\in[0,T]\,,$		(2)
	$\displaystyle u(\boldsymbol{x},0)=h(\boldsymbol{x})\,,\ \ \boldsymbol{x}\in\Omega\,,$		(3)

where $\boldsymbol{x}\in\Omega$ is a spatial vector variable in a domain $\Omega\subset R^{d}$, $t$ is time, and $\mathcal{N}_{\boldsymbol{x},t}$ is a spatial-temporal differential operator. Following Raissi et al. (2019), let $u(\boldsymbol{x},t)$ be approximated by the output $u(\boldsymbol{x},t;\boldsymbol{w})$ of a deep neural network with inputs $\boldsymbol{x}$ and $t$. Define the residual network $r(\boldsymbol{x},t;\boldsymbol{w})$, which share the same network weights $\boldsymbol{w}$ as the approximation network $u(\boldsymbol{x},t;\boldsymbol{w})$, and satisfies:

$$r(\boldsymbol{x},t;\boldsymbol{w}):=\mathcal{N}_{\boldsymbol{x},t}[u(\boldsymbol{x},t;\boldsymbol{w})]\,,$$

(4)

where all partial derivatives can be computed by automatic differentiation methods (Baydin et al., 2017; Paszke et al., 2017). The shared network weights $\boldsymbol{w}$ are trained by minimizing a loss function that penalizes the output for not satisfying (1)-(3):

$$\mathcal{L}(\boldsymbol{w})=\mathcal{L}_{s}(\boldsymbol{w})+\mathcal{L}_{r}(\boldsymbol{w})+\mathcal{L}_{b}(\boldsymbol{w})+\mathcal{L}_{0}(\boldsymbol{w})\,,$$

(5)

where $\mathcal{L}_{s}$ is the loss corresponding to sample data (if any), $\mathcal{L}_{r}$ is the loss corresponding to the residual (4), $\mathcal{L}_{b}$ is the loss due to the boundary conditions (2), and $\mathcal{L}_{0}$ is the loss due to the initial conditions (3):

	$\displaystyle\mathcal{L}_{s}(\boldsymbol{w})$	$\displaystyle=\frac{1}{N_{s}}\sum^{N_{s}}_{i=1}|u(\boldsymbol{x}^{i}_{s},t^{i}_{s};\boldsymbol{w})-y^{i}_{s}|^{2},$		(6)
	$\displaystyle\mathcal{L}_{r}(\boldsymbol{w})$	$\displaystyle=\frac{1}{N_{r}}\sum^{N_{r}}_{i=1}r(\boldsymbol{x}^{i}_{r},t^{i}_{r};\boldsymbol{w})^{2},$		(7)
	$\displaystyle\mathcal{L}_{b}(\boldsymbol{w})$	$\displaystyle=\frac{1}{N_{b}}\sum^{N_{b}}_{i=1}|u(\boldsymbol{x}^{i}_{b},t^{i}_{b};\boldsymbol{w})-g^{i}_{b}|^{2},\,$		(8)
	$\displaystyle\mathcal{L}_{0}(\boldsymbol{w})$	$\displaystyle=\frac{1}{N_{0}}\sum^{N_{0}}_{i=1}|u(\boldsymbol{x}^{i}_{0},0;\boldsymbol{w})-h_{0}^{i}|^{2}.$		(9)

where $\{\boldsymbol{x}_{s}^{i},t_{s}^{i},y_{s}^{i}\}_{i=1}^{N_{s}}$ are sample data (if any), $\{\boldsymbol{x}_{0}^{i},h_{0}^{i}=h(\boldsymbol{x}_{0}^{i})\}_{i=1}^{N_{0}}$ are initial condition point, $\{\boldsymbol{x}_{b}^{i},t^{i}_{b},g_{b}^{i}=g(\boldsymbol{x}_{b}^{i},t^{i}_{b}))\}_{i=1}^{N_{b}}$ are boundary condition points, $\{\boldsymbol{x}_{r}^{i},t^{i}_{r}\}_{i=1}^{N_{r}}$ are collocation points randomly distributed in the domain $\Omega$, and $N_{s},N_{0},N_{b}$ and $N_{r}$ denote the total number of sample data, initial points, boundary points, and collocation points, respectively. The network weights $\boldsymbol{w}$ can be tuned by minimizing the total training loss $\mathcal{L}(\boldsymbol{w})$ via standard gradient descent procedures used in deep learning.

TensorDiffEq has a boilerplate model that can loosely be followed in most instances of usage of the package. In forward problems, this process is generally described in the following order:

1. 1.

   Define the problem domain

2. 2.

   Describe the physics of the model

3. 3.

   Define the Initial Conditions and Boundary Conditions (IC/BCs)

4. 4.

   Define the neural network architecture

5. 5.

   Select and define the solver

6. 6.

   Solve the PDE using the fit method

Each of these steps has multiple options and definitions in the TensorDiffEq solution suite. The following sections will provide a brief overview of some of the built-in functionality of the package.

TensorDiffEq comes with a base class solver for inverse problems. Inverse problems can imply parameter estimation or even estimate the interactions between nonlinear operators (Lu et al., 2019) from data. TensorDiffEq contains solvers that perform parameter estimation in a PDE system. These parameters can be mobility parameters, diffusivity parameters, etc, where there is some level of a priori physical knowledge about the system in question, but a specific parameter may be unknown. TensorDiffEq contains built-in support for solving of such systems that can be solved in $N$D cases. Parameters are defined as variables that are learned over the course of the training, therefore a natural output is a trained $u(\textbf{X},t)$ solution as well as the estimate of the parameters in question.

In this article, the authors introduce TensorDiffEq, a scalable multi-GPU solver for PINNs/NeuralPDEs. Some of the main highlights of the software are covered, and more features are currently underway. TensorDiffEq contains support for various types of initial conditions, boundary conditions, and allows the user to custom-define their PDE system for their specific problem. In the event that inverse modeling is required, TensorDiffEq contains solvers that will accommodate parameter estimation of a PDE system. Currently, TensorDiffEq is the only software suite to support self-adaptive solving, demonstrated to improve training convergence and accuracy of the final solution. TensorDiffEq takes a step forward in modern implementations of PINN solvers, and fills a unique niche of being a fully open-source multi-GPU PINN solver in the current ecosystem of Scientific Machine Learning software offerings.