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

Abstract: Physics-Informed Neural Networks promise to revolutionize science and engineering practice, by introducing domain-aware deep machine learning models into scientific computation. Several software suites have emerged to make the implementation and usage of these architectures available to the research and industry communities. Here we introduce\linebreak TensorDiffEq, built on Tensorflow 2.x, which presents an intuitive Keras-like interface for problem domain definition, model definition, and solution of forward and inverse problems using physics-aware deep learning methods. TensorDiffEq takes full advantage of Tensorflow 2.x infrastructure for deployment on multiple GPUs, allowing the implementation of large high-dimensional and complex models. Simultaneously, TensorDiffEq supports the Keras API for custom neural network architecture definitions. In the case of smaller or simpler models, the package allows for rapid deployment on smaller-scale CPU platforms with negligible changes to the implementation scripts. We demonstrate the basic usage and capabilities of TensorDiffEq in solving forward, inverse, and data assimilation problems of varying sizes and levels of complexity. The source code is available at https://github.com/tensordiffeq.