nnde: A Python package for solving differential equations using neural networks

The nnde package provides a pure-Python implementation of one of the earliest approaches to using neural networks to solve differential equations the trial function method (Lagaris et al., 1998). The nnde package was initially developed as a vehicle for understanding the internal workings of feedforward neural networks, without the constraints imposed by an existing neural network framework. It has since been enhanced to provide the capability to solve differential equations of scientific interest, such as the diffusion equation described here. The ultimate goal of the package is to provide the capability to solve systems of coupled partial differential equations, such as the equations of magnetohydrodynamics.


Statement of need
The nnde package provides a pure-Python implementation of one of the earliest approaches to using neural networks to solve differential equations -the trial function method (Lagaris et al., 1998). The nnde package was initially developed as a vehicle for understanding the internal workings of feedforward neural networks, without the constraints imposed by an existing neural network framework. It has since been enhanced to provide the capability to solve differential equations of scientific interest, such as the diffusion equation described here. The ultimate goal of the package is to provide the capability to solve systems of coupled partial differential equations, such as the equations of magnetohydrodynamics.
Development of the nnde package began before the widespread adoption of modern neural network software frameworks. In the Python ecosystem, the most popular packages are TensorFlow (https://tensorflow.org) and PyTorch (https://pytorch.org). These frameworks are designed to be application-neutral -they can be used to develop neural networks with arbitrary architectures for arbitrary learning objectives. The primary advantages of these frameworks are autodifferentiation and distributed computing. By recording the sequence of mathematical operations performed in the forward pass through the network, autodifferentiation can automatically compute the gradients of the loss function with respect to each of the network parameters, as well as the network inputs. The latter capability is central to solving differential equations. Autodifferentiation also greatly reduces the volume of code that must be developed to solve a given problem. The distributed computing capability allows a network to take advantage of GPU-enabled hardware, and multiple compute nodes, to speed the calculation, with little or no code changes required. The nnde package uses a more direct method -precomputed derivative functions for the components of the differential equations of interest and the trial solution. This code is typically faster than TensorFlow or PyTorch, but requires more hand-crafted code to solve a given problem.
The most commonly used methods for solving differential equations are the Finite Element Method (FEM) and Finite Difference Method (FDM). However, these methods can be difficult to parallelize due to the need for communication between computational elements at the boundaries of the allocated subgrids. These models can also have large storage requirements for model outputs. The neural network method is straightforward to parallelize due to the independent characteristics of the computational nodes in each network layer. Additionally, the trained network solution is more compact than an FDM or FEM solution because storage of only the network weights and biases is required. The neural network solution is meshfree and does not require interpolation to retrieve the solution at a non-grid point, as is the case with FDM or FEM. Once the network is trained, computing a solution at any spatial or temporal scale requires only a series of matrix multiplications, one per network layer. The trained solution is a sum of arbitrary differentiable basis functions, and therefore the trained solution is also differentiable, which is particularly useful when computing derived quantities such as gradients and fluxes. This approach has led to several different classes of methods for solving ODEs and PDEs with neural networks. The recent surge in interest in "physicsinformed neural networks" (Raissi et al., 2019) is an indication of the dynamic nature of the field.

Description
nnde implements a version of the trial function algorithm described by Lagaris et al. (1998). This software also incorporates a modification of the trial function algorithm to automatically incorporate arbitrary Dirichlet boundary conditions of the problem directly into the neural network solution.
As a concrete example of the sort of problem that can be solved using nnde, consider the diffusion equation in two dimensions: With all boundaries fixed at 0 and with an initial condition of ψ(x, y, 0) = sin(πx) sin(πy) the analytical solution is ψ a (x, y, t) = e −2π 2 Dt sin(πx) sin(πy) The nnde package was used to create a neural network with a single hidden layer and 10 hidden nodes and trained to solve this problem. The error in the trained solution for the case of D = 0.1 is shown as a function of time in Figure 1.